polars-df 0.19.0 → 0.21.0

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 {

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

@@ -0,0 +1,14 @@
+use polars::prelude::DataTypeExpr;
+
+#[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 }
+    }
+}

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

@@ -1,15 +1,16 @@
 use std::ops::Neg;
 
-use magnus::{prelude::*, value::Opaque, IntoValue, RArray, Ruby, Value};
+use magnus::{RArray, Value};
 use polars::lazy::dsl;
 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,10 +323,15 @@ impl RbExpr {
         self.inner.clone().arg_min().into()
     }
 
-    pub fn search_sorted(&self, element: &Self, side: Wrap<SearchSortedSide>) -> Self {
+    pub fn search_sorted(
+        &self,
+        element: &Self,
+        side: Wrap<SearchSortedSide>,
+        descending: bool,
+    ) -> Self {
         self.inner
             .clone()
-            .search_sorted(element.inner.clone(), side.0)
+            .search_sorted(element.inner.clone(), side.0, descending)
             .into()
     }
 
@@ -389,16 +386,8 @@ impl RbExpr {
         strategy: String,
         limit: FillNullLimit,
     ) -> RbResult<Self> {
-        let strat = parse_fill_null_strategy(&strategy, limit)?;
-        Ok(self
-            .inner
-            .clone()
-            .apply(
-                move |s| s.fill_null(strat).map(Some),
-                GetOutput::same_type(),
-            )
-            .with_fmt("fill_null_with_strategy")
-            .into())
+        let strategy = parse_fill_null_strategy(&strategy, limit)?;
+        Ok(self.inner.clone().fill_null_with_strategy(strategy).into())
     }
 
     pub fn fill_nan(&self, expr: &Self) -> Self {
@@ -658,10 +647,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()
     }
@@ -678,10 +663,10 @@ impl RbExpr {
         self.inner.clone().upper_bound().into()
     }
 
-    pub fn cumulative_eval(&self, expr: &Self, min_periods: usize, parallel: bool) -> Self {
+    pub fn cumulative_eval(&self, expr: &Self, min_samples: usize) -> Self {
         self.inner
             .clone()
-            .cumulative_eval(expr.inner.clone(), min_periods, parallel)
+            .cumulative_eval(expr.inner.clone(), min_samples)
             .into()
     }
 
@@ -807,20 +792,10 @@ impl RbExpr {
         self.inner.clone().ewm_var(options).into()
     }
 
-    pub fn extend_constant(&self, value: Wrap<AnyValue>, n: usize) -> Self {
-        let value = value.into_value();
-        let value = Opaque::from(value);
+    pub fn extend_constant(&self, value: &Self, n: &Self) -> Self {
         self.inner
             .clone()
-            .apply(
-                move |s| {
-                    let value = Ruby::get().unwrap().get_inner(value);
-                    let value = Wrap::<AnyValue>::try_convert(value).unwrap().0;
-                    s.extend_constant(value, n).map(Some)
-                },
-                GetOutput::same_type(),
-            )
-            .with_fmt("extend")
+            .extend_constant(value.inner.clone(), n.inner.clone())
             .into()
     }
 
@@ -881,4 +856,21 @@ impl RbExpr {
             )
             .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()
+    }
 }

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

@@ -1,4 +1,4 @@
-use magnus::{prelude::*, value::Opaque, Ruby, Value};
+use magnus::{Ruby, Value, prelude::*, value::Opaque};
 use polars::lazy::dsl::lit;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
@@ -43,11 +43,15 @@ 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_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,12 +80,7 @@ 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 {
@@ -116,20 +115,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 {
@@ -198,7 +197,7 @@ impl RbExpr {
             .inner
             .clone()
             .list()
-            .to_struct(ListToStructArgs::InferWidth {
+            .to_struct(ListToStruct::InferWidth {
                 infer_field_strategy: width_strat.0,
                 get_index_name: name_gen,
                 max_fields: upper_bound,
@@ -215,4 +214,15 @@ impl RbExpr {
             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,14 +1,21 @@
 use magnus::RArray;
+use polars::prelude::Schema;
 
-use crate::{RbExpr, RbPolarsErr, RbResult};
+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)?;
+    pub fn meta_pop(&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(RArray::from_iter(
             exprs.iter().map(|e| RbExpr::from(e.clone())),
         ))
@@ -50,51 +57,21 @@ impl RbExpr {
         self.inner.clone().meta().is_regex_projection()
     }
 
-    pub fn _meta_selector_add(&self, other: &RbExpr) -> RbResult<RbExpr> {
-        let out = self
-            .inner
-            .clone()
-            .meta()
-            ._selector_add(other.inner.clone())
-            .map_err(RbPolarsErr::from)?;
-        Ok(out.into())
-    }
-
-    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_selector_and(&self, other: &RbExpr) -> RbResult<RbExpr> {
-        let out = self
-            .inner
-            .clone()
-            .meta()
-            ._selector_and(other.inner.clone())
-            .map_err(RbPolarsErr::from)?;
-        Ok(out.into())
-    }
-
-    pub fn _meta_as_selector(&self) -> RbExpr {
-        self.inner.clone().meta()._into_selector().into()
-    }
-
-    fn compute_tree_format(&self, display_as_dot: bool) -> RbResult<String> {
+    fn compute_tree_format(
+        &self,
+        display_as_dot: bool,
+        schema: Option<Wrap<Schema>>,
+    ) -> RbResult<String> {
         let e = self
             .inner
             .clone()
             .meta()
-            .into_tree_formatter(display_as_dot)
+            .into_tree_formatter(display_as_dot, schema.as_ref().map(|s| &s.0))
             .map_err(RbPolarsErr::from)?;
         Ok(format!("{e}"))
     }
 
-    pub fn meta_tree_format(&self) -> RbResult<String> {
-        self.compute_tree_format(false)
+    pub fn meta_tree_format(&self, schema: Option<Wrap<Schema>>) -> RbResult<String> {
+        self.compute_tree_format(false, schema)
     }
 }

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

@@ -1,16 +1,18 @@
 mod array;
 mod binary;
 mod categorical;
+pub mod datatype;
 mod datetime;
 mod general;
 mod list;
 mod meta;
 mod name;
 mod rolling;
+pub mod selector;
 mod string;
 mod r#struct;
 
-use magnus::{prelude::*, RArray};
+use magnus::{RArray, prelude::*};
 use polars::lazy::dsl::Expr;
 
 use crate::RbResult;

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

@@ -1,4 +1,4 @@
-use magnus::{block::Proc, value::Opaque, Ruby};
+use magnus::{Ruby, block::Proc, value::Opaque};
 use polars::prelude::*;
 use polars_utils::format_pl_smallstr;
 
@@ -20,7 +20,7 @@ impl RbExpr {
                 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(),
+                        format!("Ruby function in 'name.map' produced an error: {e}.").into(),
                     )),
                 }
             })

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

@@ -1,7 +1,7 @@
 use polars::prelude::*;
 
-use crate::conversion::Wrap;
 use crate::RbExpr;
+use crate::conversion::Wrap;
 
 impl RbExpr {
     pub fn rolling_sum(

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

@@ -0,0 +1,219 @@
+use std::hash::{Hash, Hasher};
+use std::sync::Arc;
+
+use polars::prelude::{
+    DataType, DataTypeSelector, Selector, TimeUnit, TimeUnitSet, TimeZone, TimeZoneSet,
+};
+use polars_plan::dsl;
+
+use crate::prelude::Wrap;
+use crate::{RbResult, RbTypeError};
+
+#[magnus::wrap(class = "Polars::RbSelector")]
+#[repr(transparent)]
+#[derive(Clone)]
+pub struct RbSelector {
+    pub inner: Selector,
+}
+
+impl From<Selector> for RbSelector {
+    fn from(inner: Selector) -> Self {
+        Self { inner }
+    }
+}
+
+fn parse_time_unit_set(time_units: Vec<Wrap<TimeUnit>>) -> TimeUnitSet {
+    let mut tu = TimeUnitSet::empty();
+    for v in time_units {
+        match v.0 {
+            TimeUnit::Nanoseconds => tu |= TimeUnitSet::NANO_SECONDS,
+            TimeUnit::Microseconds => tu |= TimeUnitSet::MICRO_SECONDS,
+            TimeUnit::Milliseconds => tu |= TimeUnitSet::MILLI_SECONDS,
+        }
+    }
+    tu
+}
+
+pub fn parse_datatype_selector(selector: &RbSelector) -> RbResult<DataTypeSelector> {
+    selector.inner.clone().to_dtype_selector().ok_or_else(|| {
+        RbTypeError::new_err(format!(
+            "expected datatype based expression got '{}'",
+            selector.inner
+        ))
+    })
+}
+
+impl RbSelector {
+    pub fn union(&self, other: &Self) -> Self {
+        Self {
+            inner: self.inner.clone() | other.inner.clone(),
+        }
+    }
+
+    pub fn difference(&self, other: &Self) -> Self {
+        Self {
+            inner: self.inner.clone() - other.inner.clone(),
+        }
+    }
+
+    pub fn exclusive_or(&self, other: &Self) -> Self {
+        Self {
+            inner: self.inner.clone() ^ other.inner.clone(),
+        }
+    }
+
+    pub fn intersect(&self, other: &Self) -> Self {
+        Self {
+            inner: self.inner.clone() & other.inner.clone(),
+        }
+    }
+
+    pub fn by_dtype(dtypes: Vec<Wrap<DataType>>) -> Self {
+        let dtypes = dtypes.into_iter().map(|x| x.0).collect::<Vec<_>>();
+        dsl::dtype_cols(dtypes).as_selector().into()
+    }
+
+    pub fn by_name(names: Vec<String>, strict: bool) -> Self {
+        dsl::by_name(names, strict).into()
+    }
+
+    pub fn by_index(indices: Vec<i64>, strict: bool) -> Self {
+        Selector::ByIndex {
+            indices: indices.into(),
+            strict,
+        }
+        .into()
+    }
+
+    pub fn first(strict: bool) -> Self {
+        Selector::ByIndex {
+            indices: [0].into(),
+            strict,
+        }
+        .into()
+    }
+
+    pub fn last(strict: bool) -> Self {
+        Selector::ByIndex {
+            indices: [-1].into(),
+            strict,
+        }
+        .into()
+    }
+
+    pub fn matches(pattern: String) -> Self {
+        Selector::Matches(pattern.into()).into()
+    }
+
+    pub fn enum_() -> Self {
+        DataTypeSelector::Enum.as_selector().into()
+    }
+
+    pub fn categorical() -> Self {
+        DataTypeSelector::Categorical.as_selector().into()
+    }
+
+    pub fn nested() -> Self {
+        DataTypeSelector::Nested.as_selector().into()
+    }
+
+    pub fn list(inner_dst: Option<&Self>) -> RbResult<Self> {
+        let inner_dst = match inner_dst {
+            None => None,
+            Some(inner_dst) => Some(Arc::new(parse_datatype_selector(inner_dst)?)),
+        };
+        Ok(DataTypeSelector::List(inner_dst).as_selector().into())
+    }
+
+    pub fn array(inner_dst: Option<&Self>, width: Option<usize>) -> RbResult<Self> {
+        let inner_dst = match inner_dst {
+            None => None,
+            Some(inner_dst) => Some(Arc::new(parse_datatype_selector(inner_dst)?)),
+        };
+        Ok(DataTypeSelector::Array(inner_dst, width)
+            .as_selector()
+            .into())
+    }
+
+    pub fn struct_() -> Self {
+        DataTypeSelector::Struct.as_selector().into()
+    }
+
+    pub fn integer() -> Self {
+        DataTypeSelector::Integer.as_selector().into()
+    }
+
+    pub fn signed_integer() -> Self {
+        DataTypeSelector::SignedInteger.as_selector().into()
+    }
+
+    pub fn unsigned_integer() -> Self {
+        DataTypeSelector::UnsignedInteger.as_selector().into()
+    }
+
+    pub fn float() -> Self {
+        DataTypeSelector::Float.as_selector().into()
+    }
+
+    pub fn decimal() -> Self {
+        DataTypeSelector::Decimal.as_selector().into()
+    }
+
+    pub fn numeric() -> Self {
+        DataTypeSelector::Numeric.as_selector().into()
+    }
+
+    pub fn temporal() -> Self {
+        DataTypeSelector::Temporal.as_selector().into()
+    }
+
+    pub fn datetime(tu: Vec<Wrap<TimeUnit>>, tz: Vec<Wrap<Option<TimeZone>>>) -> Self {
+        use TimeZoneSet as TZS;
+
+        let mut allow_unset = false;
+        let mut allow_set = false;
+        let mut any_of: Vec<TimeZone> = Vec::new();
+
+        let tu = parse_time_unit_set(tu);
+        for t in tz {
+            let t = t.0;
+            match t {
+                None => allow_unset = true,
+                Some(s) if s.as_str() == "*" => allow_set = true,
+                Some(t) => any_of.push(t),
+            }
+        }
+
+        let tzs = match (allow_unset, allow_set) {
+            (true, true) => TZS::Any,
+            (false, true) => TZS::AnySet,
+            (true, false) if any_of.is_empty() => TZS::Unset,
+            (true, false) => TZS::UnsetOrAnyOf(any_of.into()),
+            (false, false) => TZS::AnyOf(any_of.into()),
+        };
+        DataTypeSelector::Datetime(tu, tzs).as_selector().into()
+    }
+
+    pub fn duration(tu: Vec<Wrap<TimeUnit>>) -> Self {
+        let tu = parse_time_unit_set(tu);
+        DataTypeSelector::Duration(tu).as_selector().into()
+    }
+
+    pub fn object() -> Self {
+        DataTypeSelector::Object.as_selector().into()
+    }
+
+    pub fn empty() -> Self {
+        dsl::empty().into()
+    }
+
+    pub fn all() -> Self {
+        dsl::all().into()
+    }
+
+    pub fn hash(&self) -> u64 {
+        let mut hasher = std::hash::DefaultHasher::default();
+        self.inner.hash(&mut hasher);
+        hasher.finish()
+    }
+}

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

@@ -163,12 +163,20 @@ impl RbExpr {
         self.inner.clone().str().reverse().into()
     }
 
-    pub fn str_pad_start(&self, length: usize, fillchar: char) -> Self {
-        self.clone().inner.str().pad_start(length, fillchar).into()
+    pub fn str_pad_start(&self, length: &RbExpr, fillchar: char) -> Self {
+        self.clone()
+            .inner
+            .str()
+            .pad_start(length.inner.clone(), fillchar)
+            .into()
     }
 
-    pub fn str_pad_end(&self, length: usize, fillchar: char) -> Self {
-        self.clone().inner.str().pad_end(length, fillchar).into()
+    pub fn str_pad_end(&self, length: &RbExpr, fillchar: char) -> Self {
+        self.clone()
+            .inner
+            .str()
+            .pad_end(length.inner.clone(), fillchar)
+            .into()
     }
 
     pub fn str_zfill(&self, length: &Self) -> Self {
@@ -220,12 +228,11 @@ impl RbExpr {
         self.inner.clone().str().base64_decode(strict).into()
     }
 
-    pub fn str_to_integer(&self, base: &Self, strict: bool) -> Self {
+    pub fn str_to_integer(&self, base: &Self, dtype: Option<Wrap<DataType>>, strict: bool) -> Self {
         self.inner
             .clone()
             .str()
-            .to_integer(base.inner.clone(), strict)
-            .with_fmt("str.to_integer")
+            .to_integer(base.inner.clone(), dtype.map(|wrap| wrap.0), strict)
             .into()
     }