polars-df 0.15.0 → 0.26.0

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

@@ -1,38 +1,45 @@
+use magnus::Value;
 use polars::prelude::*;
 
-use crate::RbExpr;
+use crate::ruby::plan_callback::PlanCallbackExt;
+use crate::ruby::ruby_function::RubyObject;
+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_unique(&self, maintain_order: bool) -> Self {
-        if maintain_order {
-            self.inner.clone().arr().unique_stable().into()
-        } else {
-            self.inner.clone().arr().unique().into()
-        }
+    pub fn arr_std(&self, ddof: u8) -> Self {
+        self.inner.clone().arr().std(ddof).into()
     }
 
-    pub fn arr_to_list(&self) -> Self {
-        self.inner.clone().arr().to_list().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_all(&self) -> Self {
-        self.inner.clone().arr().all().into()
+    pub fn arr_median(&self) -> Self {
+        self.inner.clone().arr().median().into()
     }
 
-    pub fn arr_any(&self) -> Self {
-        self.inner.clone().arr().any().into()
+    pub fn arr_to_list(&self) -> Self {
+        self.inner.clone().arr().to_list().into()
     }
 
     pub fn arr_sort(&self, descending: bool, nulls_last: bool) -> Self {
@@ -47,10 +54,6 @@ impl RbExpr {
             .into()
     }
 
-    pub fn arr_reverse(&self) -> Self {
-        self.inner.clone().arr().reverse().into()
-    }
-
     pub fn arr_arg_min(&self) -> Self {
         self.inner.clone().arr().arg_min().into()
     }
@@ -75,11 +78,11 @@ impl RbExpr {
             .into()
     }
 
-    pub fn arr_contains(&self, other: &RbExpr) -> Self {
+    pub fn arr_contains(&self, other: &RbExpr, nulls_equal: bool) -> Self {
         self.inner
             .clone()
             .arr()
-            .contains(other.inner.clone())
+            .contains(other.inner.clone(), nulls_equal)
             .into()
     }
 
@@ -90,4 +93,65 @@ impl RbExpr {
             .count_matches(expr.inner.clone())
             .into()
     }
+
+    pub fn arr_to_struct(&self, name_gen: Option<Value>) -> Self {
+        let name_gen = name_gen.map(|o| PlanCallback::new_ruby(RubyObject::from(o)));
+        self.inner.clone().arr().to_struct(name_gen).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, empty_as_null: bool, keep_nulls: bool) -> Self {
+        self.inner
+            .clone()
+            .arr()
+            .explode(ExplodeOptions {
+                empty_as_null,
+                keep_nulls,
+            })
+            .into()
+    }
+
+    pub fn arr_eval(&self, expr: &RbExpr, as_list: bool) -> Self {
+        self.inner
+            .clone()
+            .arr()
+            .eval(expr.inner.clone(), as_list)
+            .into()
+    }
+
+    pub fn arr_agg(&self, expr: &RbExpr) -> Self {
+        self.inner.clone().arr().agg(expr.inner.clone()).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,52 @@ 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()
+    }
+
+    pub fn bin_slice(&self, offset: &RbExpr, length: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .binary()
+            .slice(offset.inner.clone(), length.inner.clone())
+            .into()
+    }
+
+    pub fn bin_head(&self, n: &RbExpr) -> Self {
+        self.inner.clone().binary().head(n.inner.clone()).into()
+    }
+
+    pub fn bin_tail(&self, n: &RbExpr) -> Self {
+        self.inner.clone().binary().tail(n.inner.clone()).into()
+    }
+
+    pub fn bin_get(&self, index: &RbExpr, null_on_oob: bool) -> Self {
+        self.inner
+            .clone()
+            .binary()
+            .get(index.inner.clone(), null_on_oob)
+            .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,51 @@
+use magnus::{RArray, Ruby, TryConvert, Value};
+use polars::prelude::{DataType, DataTypeExpr, PlSmallStr, Schema};
+
+use crate::prelude::Wrap;
+use crate::ruby::utils::TryIntoValue;
+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 self_dtype() -> Self {
+        DataTypeExpr::SelfDtype.into()
+    }
+
+    pub fn collect_dtype(ruby: &Ruby, self_: &Self, schema: Wrap<Schema>) -> RbResult<Value> {
+        let dtype = self_
+            .clone()
+            .inner
+            .into_datatype(&schema.0)
+            .map_err(RbPolarsErr::from)?;
+        Wrap(dtype).try_into_value_with(ruby)
+    }
+
+    pub fn struct_with_fields(rb_fields: RArray) -> RbResult<Self> {
+        let mut fields = Vec::new();
+        for v in rb_fields.into_iter() {
+            let (name, dt_expr) = <(String, &RbDataTypeExpr)>::try_convert(v)?;
+            fields.push((PlSmallStr::from_string(name), dt_expr.inner.clone()));
+        }
+        Ok(DataTypeExpr::StructWithFields(fields).into())
+    }
+}

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

@@ -1,9 +1,23 @@
 use polars::prelude::*;
 
 use crate::conversion::Wrap;
-use crate::RbExpr;
+use crate::{RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
+    pub fn dt_add_business_days(
+        &self,
+        n: &RbExpr,
+        week_mask: [bool; 7],
+        holidays: &RbExpr,
+        roll: Wrap<Roll>,
+    ) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .add_business_days(n.inner.clone(), week_mask, holidays.inner.clone(), roll.0)
+            .into()
+    }
+
     pub fn dt_to_string(&self, format: String) -> Self {
         self.inner.clone().dt().to_string(&format).into()
     }
@@ -12,30 +26,17 @@ 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()
-    }
-
-    pub fn dt_convert_time_zone(&self, time_zone: String) -> Self {
-        self.inner
+    pub fn dt_convert_time_zone(&self, time_zone: String) -> RbResult<Self> {
+        Ok(self
+            .inner
             .clone()
             .dt()
-            .convert_time_zone(time_zone.into())
-            .into()
+            .convert_time_zone(
+                TimeZone::opt_try_new(Some(PlSmallStr::from(time_zone)))
+                    .map_err(RbPolarsErr::from)?
+                    .unwrap_or(TimeZone::UTC),
+            )
+            .into())
     }
 
     pub fn dt_cast_time_unit(&self, tu: Wrap<TimeUnit>) -> Self {
@@ -47,16 +48,18 @@ impl RbExpr {
         time_zone: Option<String>,
         ambiguous: &Self,
         non_existent: Wrap<NonExistent>,
-    ) -> Self {
-        self.inner
+    ) -> RbResult<Self> {
+        Ok(self
+            .inner
             .clone()
             .dt()
             .replace_time_zone(
-                time_zone.map(|x| x.into()),
+                TimeZone::opt_try_new(time_zone.map(PlSmallStr::from_string))
+                    .map_err(RbPolarsErr::from)?,
                 ambiguous.inner.clone(),
                 non_existent.0,
             )
-            .into()
+            .into())
     }
 
     pub fn dt_truncate(&self, every: &Self) -> Self {
@@ -83,6 +86,33 @@ impl RbExpr {
         self.inner.clone().dt().round(every.inner.clone()).into()
     }
 
+    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()
@@ -91,10 +121,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: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .is_business_day(week_mask, holidays.inner.clone())
+            .into()
+    }
+
     pub fn dt_is_leap_year(&self) -> Self {
         self.clone().inner.dt().is_leap_year().into()
     }
@@ -111,6 +157,10 @@ impl RbExpr {
         self.clone().inner.dt().month().into()
     }
 
+    pub fn dt_days_in_month(&self) -> Self {
+        self.inner.clone().dt().days_in_month().into()
+    }
+
     pub fn dt_week(&self) -> Self {
         self.clone().inner.dt().week().into()
     }
@@ -167,31 +217,39 @@ impl RbExpr {
         self.inner.clone().dt().timestamp(tu.0).into()
     }
 
-    pub fn dt_total_days(&self) -> Self {
-        self.inner.clone().dt().total_days().into()
+    pub fn dt_total_days(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_days(fractional).into()
     }
 
-    pub fn dt_total_hours(&self) -> Self {
-        self.inner.clone().dt().total_hours().into()
+    pub fn dt_total_hours(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_hours(fractional).into()
     }
 
-    pub fn dt_total_minutes(&self) -> Self {
-        self.inner.clone().dt().total_minutes().into()
+    pub fn dt_total_minutes(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_minutes(fractional).into()
     }
 
-    pub fn dt_total_seconds(&self) -> Self {
-        self.inner.clone().dt().total_seconds().into()
+    pub fn dt_total_seconds(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_seconds(fractional).into()
     }
 
-    pub fn dt_total_milliseconds(&self) -> Self {
-        self.inner.clone().dt().total_milliseconds().into()
+    pub fn dt_total_milliseconds(&self, fractional: bool) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .total_milliseconds(fractional)
+            .into()
     }
 
-    pub fn dt_total_microseconds(&self) -> Self {
-        self.inner.clone().dt().total_microseconds().into()
+    pub fn dt_total_microseconds(&self, fractional: bool) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .total_microseconds(fractional)
+            .into()
     }
 
-    pub fn dt_total_nanoseconds(&self) -> Self {
-        self.inner.clone().dt().total_nanoseconds().into()
+    pub fn dt_total_nanoseconds(&self, fractional: bool) -> Self {
+        self.inner.clone().dt().total_nanoseconds(fractional).into()
     }
 }

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

@@ -0,0 +1,12 @@
+use crate::RbExpr;
+use crate::expr::datatype::RbDataTypeExpr;
+
+impl RbExpr {
+    pub fn ext_storage(&self) -> Self {
+        self.inner.clone().ext().storage().into()
+    }
+
+    pub fn ext_to(&self, dtype: &RbDataTypeExpr) -> Self {
+        self.inner.clone().ext().to(dtype.inner.clone()).into()
+    }
+}