polars-df 0.15.0 → 0.26.0

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

@@ -1,4 +1,4 @@
-use magnus::{prelude::*, value::Opaque, Ruby, Value};
+use magnus::{RArray, prelude::*};
 use polars::lazy::dsl::lit;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
@@ -7,14 +7,6 @@ use crate::conversion::Wrap;
 use crate::{RbExpr, RbResult};
 
 impl RbExpr {
-    pub fn list_all(&self) -> Self {
-        self.inner.clone().list().all().into()
-    }
-
-    pub fn list_any(&self) -> Self {
-        self.inner.clone().list().any().into()
-    }
-
     pub fn list_arg_max(&self) -> Self {
         self.inner.clone().list().arg_max().into()
     }
@@ -23,11 +15,11 @@ impl RbExpr {
         self.inner.clone().list().arg_min().into()
     }
 
-    pub fn list_contains(&self, other: &RbExpr) -> Self {
+    pub fn list_contains(&self, other: &RbExpr, nulls_equal: bool) -> Self {
         self.inner
             .clone()
             .list()
-            .contains(other.inner.clone())
+            .contains(other.inner.clone(), nulls_equal)
             .into()
     }
 
@@ -43,11 +35,19 @@ impl RbExpr {
         Ok(self.inner.clone().list().diff(n, null_behavior.0).into())
     }
 
-    pub fn list_eval(&self, expr: &RbExpr, parallel: bool) -> Self {
+    pub fn list_eval(&self, expr: &RbExpr) -> Self {
+        self.inner.clone().list().eval(expr.inner.clone()).into()
+    }
+
+    pub fn list_agg(&self, expr: &RbExpr) -> Self {
+        self.inner.clone().list().agg(expr.inner.clone()).into()
+    }
+
+    pub fn list_filter(&self, predicate: &RbExpr) -> Self {
         self.inner
             .clone()
             .list()
-            .eval(expr.inner.clone(), parallel)
+            .eval(Expr::Column(PlSmallStr::EMPTY).filter(predicate.inner.clone()))
             .into()
     }
 
@@ -76,20 +76,23 @@ impl RbExpr {
     }
 
     pub fn list_mean(&self) -> Self {
-        self.inner
-            .clone()
-            .list()
-            .mean()
-            .with_fmt("list.mean")
-            .into()
+        self.inner.clone().list().mean().into()
     }
 
-    pub fn list_min(&self) -> Self {
-        self.inner.clone().list().min().into()
+    pub fn list_median(&self) -> Self {
+        self.inner.clone().list().median().into()
     }
 
-    pub fn list_reverse(&self) -> Self {
-        self.inner.clone().list().reverse().into()
+    pub fn list_std(&self, ddof: u8) -> Self {
+        self.inner.clone().list().std(ddof).into()
+    }
+
+    pub fn list_var(&self, ddof: u8) -> Self {
+        self.inner.clone().list().var(ddof).into()
+    }
+
+    pub fn list_min(&self) -> Self {
+        self.inner.clone().list().min().into()
     }
 
     pub fn list_shift(&self, periods: &RbExpr) -> Self {
@@ -116,20 +119,20 @@ impl RbExpr {
         self.inner.clone().list().tail(n.inner.clone()).into()
     }
 
-    pub fn list_sort(&self, reverse: bool) -> Self {
+    pub fn list_sort(&self, descending: bool, nulls_last: bool) -> Self {
         self.inner
             .clone()
             .list()
-            .sort(SortOptions {
-                descending: reverse,
-                ..Default::default()
-            })
-            .with_fmt("list.sort")
+            .sort(
+                SortOptions::default()
+                    .with_order_descending(descending)
+                    .with_nulls_last(nulls_last),
+            )
             .into()
     }
 
     pub fn list_sum(&self) -> Self {
-        self.inner.clone().list().sum().with_fmt("list.sum").into()
+        self.inner.clone().list().sum().into()
     }
 
     pub fn list_drop_nulls(&self) -> Self {
@@ -172,47 +175,40 @@ impl RbExpr {
             .into()
     }
 
+    pub fn list_gather_every(&self, n: &RbExpr, offset: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .gather_every(n.inner.clone(), offset.inner.clone())
+            .into()
+    }
+
     pub fn list_to_array(&self, width: usize) -> Self {
         self.inner.clone().list().to_array(width).into()
     }
 
-    pub fn list_to_struct(
-        &self,
-        width_strat: Wrap<ListToStructWidthStrategy>,
-        name_gen: Option<Value>,
-        upper_bound: usize,
-    ) -> RbResult<Self> {
-        let name_gen = name_gen.map(|lambda| {
-            let lambda = Opaque::from(lambda);
-            Arc::new(move |idx: usize| {
-                let lambda = Ruby::get().unwrap().get_inner(lambda);
-                let out: String = lambda.funcall("call", (idx,)).unwrap();
-                PlSmallStr::from_string(out)
-            });
-
-            // non-Ruby thread
-            todo!();
-        });
-
+    pub fn list_to_struct(&self, names: RArray) -> RbResult<Self> {
         Ok(self
             .inner
             .clone()
             .list()
-            .to_struct(ListToStructArgs::InferWidth {
-                infer_field_strategy: width_strat.0,
-                get_index_name: name_gen,
-                max_fields: upper_bound,
-            })
+            .to_struct(
+                names
+                    .into_iter()
+                    .map(|x| Ok(Wrap::<PlSmallStr>::try_convert(x)?.0))
+                    .collect::<RbResult<Arc<[_]>>>()?,
+            )
             .into())
     }
 
-    pub fn list_unique(&self, maintain_order: bool) -> Self {
-        let e = self.inner.clone();
-
-        if maintain_order {
-            e.list().unique_stable().into()
-        } else {
-            e.list().unique().into()
+    pub fn list_set_operation(&self, other: &RbExpr, operation: Wrap<SetOperation>) -> Self {
+        let e = self.inner.clone().list();
+        match operation.0 {
+            SetOperation::Intersection => e.set_intersection(other.inner.clone()),
+            SetOperation::Difference => e.set_difference(other.inner.clone()),
+            SetOperation::Union => e.union(other.inner.clone()),
+            SetOperation::SymmetricDifference => e.set_symmetric_difference(other.inner.clone()),
         }
+        .into()
     }
 }

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

@@ -1,17 +1,23 @@
-use magnus::RArray;
+use magnus::{RArray, Ruby};
+use polars::prelude::Schema;
 
-use crate::{RbExpr, RbPolarsErr, RbResult};
+use crate::expr::ToRbExprs;
+use crate::{RbExpr, RbPolarsErr, RbResult, Wrap};
 
 impl RbExpr {
     pub fn meta_eq(&self, other: &RbExpr) -> bool {
         self.inner == other.inner
     }
 
-    pub fn meta_pop(&self) -> RbResult<RArray> {
-        let exprs = self.inner.clone().meta().pop().map_err(RbPolarsErr::from)?;
-        Ok(RArray::from_iter(
-            exprs.iter().map(|e| RbExpr::from(e.clone())),
-        ))
+    pub fn meta_pop(ruby: &Ruby, self_: &Self, schema: Option<Wrap<Schema>>) -> RbResult<RArray> {
+        let schema = schema.as_ref().map(|s| &s.0);
+        let exprs = self_
+            .inner
+            .clone()
+            .meta()
+            .pop(schema)
+            .map_err(RbPolarsErr::from)?;
+        Ok(exprs.to_rbexprs(ruby))
     }
 
     pub fn meta_root_names(&self) -> Vec<String> {
@@ -50,47 +56,36 @@ impl RbExpr {
         self.inner.clone().meta().is_regex_projection()
     }
 
-    pub fn _meta_selector_add(&self, other: &RbExpr) -> RbResult<RbExpr> {
-        let out = self
-            .inner
+    pub fn meta_is_column_selection(&self, allow_aliasing: bool) -> bool {
+        self.inner
             .clone()
             .meta()
-            ._selector_add(other.inner.clone())
-            .map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+            .is_column_selection(allow_aliasing)
     }
 
-    pub fn _meta_selector_sub(&self, other: &RbExpr) -> RbResult<RbExpr> {
-        let out = self
-            .inner
-            .clone()
-            .meta()
-            ._selector_sub(other.inner.clone())
-            .map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+    pub fn meta_is_literal(&self, allow_aliasing: bool) -> bool {
+        self.inner.clone().meta().is_literal(allow_aliasing)
     }
 
-    pub fn _meta_selector_and(&self, other: &RbExpr) -> RbResult<RbExpr> {
-        let out = self
+    fn compute_tree_format(
+        &self,
+        display_as_dot: bool,
+        schema: Option<Wrap<Schema>>,
+    ) -> RbResult<String> {
+        let e = self
             .inner
             .clone()
             .meta()
-            ._selector_and(other.inner.clone())
+            .into_tree_formatter(display_as_dot, schema.as_ref().map(|s| &s.0))
             .map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+        Ok(format!("{e}"))
     }
 
-    pub fn _meta_as_selector(&self) -> RbExpr {
-        self.inner.clone().meta()._into_selector().into()
+    pub fn meta_tree_format(&self, schema: Option<Wrap<Schema>>) -> RbResult<String> {
+        self.compute_tree_format(false, schema)
     }
 
-    pub fn meta_tree_format(&self) -> RbResult<String> {
-        let e = self
-            .inner
-            .clone()
-            .meta()
-            .into_tree_formatter()
-            .map_err(RbPolarsErr::from)?;
-        Ok(format!("{e}"))
+    pub fn meta_show_graph(&self, schema: Option<Wrap<Schema>>) -> RbResult<String> {
+        self.compute_tree_format(true, schema)
     }
 }

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

@@ -1,16 +1,22 @@
 mod array;
 mod binary;
+mod bitwise;
 mod categorical;
+pub mod datatype;
 mod datetime;
+mod extension;
 mod general;
 mod list;
 mod meta;
 mod name;
 mod rolling;
+pub mod selector;
+#[cfg(feature = "serialize_binary")]
+mod serde;
 mod string;
 mod r#struct;
 
-use magnus::{prelude::*, RArray};
+use magnus::{RArray, Ruby, prelude::*};
 use polars::lazy::dsl::Expr;
 
 use crate::RbResult;
@@ -27,10 +33,26 @@ impl From<Expr> for RbExpr {
     }
 }
 
-pub fn rb_exprs_to_exprs(rb_exprs: RArray) -> RbResult<Vec<Expr>> {
-    let mut exprs = Vec::new();
-    for item in rb_exprs.into_iter() {
-        exprs.push(<&RbExpr>::try_convert(item)?.inner.clone());
+pub(crate) trait ToExprs {
+    fn to_exprs(self) -> RbResult<Vec<Expr>>;
+}
+
+impl ToExprs for RArray {
+    fn to_exprs(self) -> RbResult<Vec<Expr>> {
+        let mut exprs = Vec::new();
+        for item in self.into_iter() {
+            exprs.push(<&RbExpr>::try_convert(item)?.inner.clone());
+        }
+        Ok(exprs)
+    }
+}
+
+pub(crate) trait ToRbExprs {
+    fn to_rbexprs(self, rb: &Ruby) -> RArray;
+}
+
+impl ToRbExprs for Vec<Expr> {
+    fn to_rbexprs(self, rb: &Ruby) -> RArray {
+        rb.ary_from_iter(self.iter().map(|e| RbExpr::from(e.clone())))
     }
-    Ok(exprs)
 }

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

@@ -1,29 +1,20 @@
-use magnus::{block::Proc, value::Opaque, Ruby};
+use magnus::Value;
 use polars::prelude::*;
-use polars_utils::format_pl_smallstr;
 
 use crate::RbExpr;
+use crate::ruby::plan_callback::PlanCallbackExt;
+use crate::ruby::ruby_function::RubyObject;
 
 impl RbExpr {
     pub fn name_keep(&self) -> Self {
         self.inner.clone().name().keep().into()
     }
 
-    pub fn name_map(&self, lambda: Proc) -> Self {
-        let lambda = Opaque::from(lambda);
+    pub fn name_map(&self, lambda: Value) -> Self {
         self.inner
             .clone()
             .name()
-            .map(move |name| {
-                let lambda = Ruby::get().unwrap().get_inner(lambda);
-                let out = lambda.call::<_, String>((name.as_str(),));
-                match out {
-                    Ok(out) => Ok(format_pl_smallstr!("{}", out)),
-                    Err(e) => Err(PolarsError::ComputeError(
-                        format!("Ruby function in 'name.map' produced an error: {}.", e).into(),
-                    )),
-                }
-            })
+            .map(PlanCallback::new_ruby(RubyObject::from(lambda)))
             .into()
     }
 
@@ -42,4 +33,28 @@ impl RbExpr {
     pub fn name_to_uppercase(&self) -> Self {
         self.inner.clone().name().to_uppercase().into()
     }
+
+    pub fn name_replace(&self, pattern: String, value: String, literal: bool) -> Self {
+        self.inner
+            .clone()
+            .name()
+            .replace(&pattern, &value, literal)
+            .into()
+    }
+
+    pub fn name_map_fields(&self, name_mapper: Value) -> Self {
+        self.inner
+            .clone()
+            .name()
+            .map_fields(PlanCallback::new_ruby(RubyObject::from(name_mapper)))
+            .into()
+    }
+
+    pub fn name_prefix_fields(&self, prefix: String) -> Self {
+        self.inner.clone().name().prefix_fields(&prefix).into()
+    }
+
+    pub fn name_suffix_fields(&self, suffix: String) -> Self {
+        self.inner.clone().name().suffix_fields(&suffix).into()
+    }
 }

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

@@ -1,7 +1,10 @@
+use magnus::Value;
 use polars::prelude::*;
 
 use crate::conversion::Wrap;
-use crate::RbExpr;
+use crate::ruby::plan_callback::PlanCallbackExt;
+use crate::ruby::ruby_function::RubyObject;
+use crate::{RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
     pub fn rolling_sum(
@@ -319,7 +322,112 @@ impl RbExpr {
             .into()
     }
 
-    pub fn rolling_skew(&self, window_size: usize, bias: bool) -> Self {
-        self.inner.clone().rolling_skew(window_size, bias).into()
+    pub fn rolling_rank(
+        &self,
+        window_size: usize,
+        method: Wrap<RollingRankMethod>,
+        seed: Option<u64>,
+        min_samples: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_samples = min_samples.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            window_size,
+            min_periods: min_samples,
+            weights: None,
+            center,
+            fn_params: Some(RollingFnParams::Rank {
+                method: method.0,
+                seed,
+            }),
+        };
+
+        self.inner.clone().rolling_rank(options).into()
+    }
+
+    pub fn rolling_rank_by(
+        &self,
+        by: &RbExpr,
+        window_size: String,
+        method: Wrap<RollingRankMethod>,
+        seed: Option<u64>,
+        min_samples: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> RbResult<Self> {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::try_parse(&window_size).map_err(RbPolarsErr::from)?,
+            min_periods: min_samples,
+            closed_window: closed.0,
+            fn_params: Some(RollingFnParams::Rank {
+                method: method.0,
+                seed,
+            }),
+        };
+
+        Ok(self
+            .inner
+            .clone()
+            .rolling_rank_by(by.inner.clone(), options)
+            .into())
+    }
+
+    pub fn rolling_skew(
+        &self,
+        window_size: usize,
+        bias: bool,
+        min_periods: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            window_size,
+            weights: None,
+            min_periods,
+            center,
+            fn_params: Some(RollingFnParams::Skew { bias }),
+        };
+
+        self.inner.clone().rolling_skew(options).into()
+    }
+
+    pub fn rolling_kurtosis(
+        &self,
+        window_size: usize,
+        fisher: bool,
+        bias: bool,
+        min_periods: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            window_size,
+            weights: None,
+            min_periods,
+            center,
+            fn_params: Some(RollingFnParams::Kurtosis { fisher, bias }),
+        };
+
+        self.inner.clone().rolling_kurtosis(options).into()
+    }
+
+    pub fn rolling_map(
+        &self,
+        lambda: Value,
+        window_size: usize,
+        weights: Option<Vec<f64>>,
+        min_periods: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            window_size,
+            weights,
+            min_periods,
+            center,
+            ..Default::default()
+        };
+        let function = PlanCallback::new_ruby(RubyObject::from(lambda));
+
+        self.inner.clone().rolling_map(function, options).into()
     }
 }