polars-df 0.15.0 → 0.26.0

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::functions::dtype_cols(dtypes).as_selector().into()
+    }
+
+    pub fn by_name(names: Vec<String>, strict: bool, expand_patterns: bool) -> Self {
+        dsl::functions::by_name(names, strict, expand_patterns).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::functions::empty().into()
+    }
+
+    pub fn all() -> Self {
+        dsl::functions::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/serde.rs

@@ -0,0 +1,28 @@
+use std::io::{BufReader, BufWriter};
+
+use magnus::Value;
+use polars::lazy::prelude::Expr;
+use polars_utils::pl_serialize;
+
+use crate::exceptions::ComputeError;
+use crate::file::get_file_like;
+use crate::{RbExpr, RbResult};
+
+impl RbExpr {
+    pub fn serialize_binary(&self, rb_f: Value) -> RbResult<()> {
+        let file = get_file_like(rb_f, true)?;
+        let writer = BufWriter::new(file);
+        pl_serialize::SerializeOptions::default()
+            .serialize_into_writer::<_, _, true>(writer, &self.inner)
+            .map_err(|err| ComputeError::new_err(err.to_string()))
+    }
+
+    pub fn deserialize_binary(rb_f: Value) -> RbResult<RbExpr> {
+        let file = get_file_like(rb_f, false)?;
+        let reader = BufReader::new(file);
+        let expr: Expr = pl_serialize::SerializeOptions::default()
+            .deserialize_from_reader::<_, _, true>(reader)
+            .map_err(|err| ComputeError::new_err(err.to_string()))?;
+        Ok(expr.into())
+    }
+}

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

@@ -1,7 +1,7 @@
 use polars::prelude::*;
 
 use crate::conversion::Wrap;
-use crate::{RbExpr, RbPolarsErr, RbResult};
+use crate::{RbDataTypeExpr, RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
     pub fn str_join(&self, delimiter: String, ignore_nulls: bool) -> Self {
@@ -30,19 +30,18 @@ impl RbExpr {
         self.inner.clone().str().to_date(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn str_to_datetime(
         &self,
         format: Option<String>,
         time_unit: Option<Wrap<TimeUnit>>,
-        time_zone: Option<Wrap<TimeZone>>,
+        time_zone: Wrap<Option<TimeZone>>,
         strict: bool,
         exact: bool,
         cache: bool,
         ambiguous: &Self,
     ) -> Self {
         let format = format.map(|x| x.into());
-        let time_zone = time_zone.map(|x| x.0);
+        let time_zone = time_zone.0;
 
         let options = StrptimeOptions {
             format,
@@ -122,6 +121,14 @@ impl RbExpr {
             .into()
     }
 
+    pub fn str_head(&self, n: &Self) -> Self {
+        self.inner.clone().str().head(n.inner.clone()).into()
+    }
+
+    pub fn str_tail(&self, n: &Self) -> Self {
+        self.inner.clone().str().tail(n.inner.clone()).into()
+    }
+
     pub fn str_to_uppercase(&self) -> Self {
         self.inner.clone().str().to_uppercase().into()
     }
@@ -159,16 +166,28 @@ impl RbExpr {
             .into()
     }
 
+    pub fn str_normalize(&self, form: Wrap<UnicodeForm>) -> Self {
+        self.inner.clone().str().normalize(form.0).into()
+    }
+
     pub fn str_reverse(&self) -> Self {
         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 {
@@ -192,6 +211,23 @@ impl RbExpr {
         }
     }
 
+    pub fn str_find(&self, pat: &Self, literal: Option<bool>, strict: bool) -> Self {
+        match literal {
+            Some(true) => self
+                .inner
+                .clone()
+                .str()
+                .find_literal(pat.inner.clone())
+                .into(),
+            _ => self
+                .inner
+                .clone()
+                .str()
+                .find(pat.inner.clone(), strict)
+                .into(),
+        }
+    }
+
     pub fn str_ends_with(&self, sub: &RbExpr) -> Self {
         self.inner.clone().str().ends_with(sub.inner.clone()).into()
     }
@@ -220,25 +256,19 @@ 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()
     }
 
-    pub fn str_json_decode(
-        &self,
-        dtype: Option<Wrap<DataType>>,
-        infer_schema_len: Option<usize>,
-    ) -> Self {
-        let dtype = dtype.map(|wrap| wrap.0);
+    pub fn str_json_decode(&self, dtype: &RbDataTypeExpr) -> Self {
         self.inner
             .clone()
             .str()
-            .json_decode(dtype, infer_schema_len)
+            .json_decode(dtype.inner.clone())
             .into()
     }
 
@@ -316,8 +346,32 @@ impl RbExpr {
         self.inner.clone().str().splitn(by.inner.clone(), n).into()
     }
 
-    pub fn str_to_decimal(&self, infer_len: usize) -> Self {
-        self.inner.clone().str().to_decimal(infer_len).into()
+    pub fn str_split_regex(&self, by: &Self, strict: bool) -> Self {
+        self.str_split_regex_with_strict(by, strict)
+    }
+
+    pub fn str_split_regex_inclusive(&self, by: &Self, strict: bool) -> Self {
+        self.str_split_regex_inclusive_with_strict(by, strict)
+    }
+
+    fn str_split_regex_with_strict(&self, by: &Self, strict: bool) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .split_regex(by.inner.clone(), strict)
+            .into()
+    }
+
+    fn str_split_regex_inclusive_with_strict(&self, by: &Self, strict: bool) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .split_regex_inclusive(by.inner.clone(), strict)
+            .into()
+    }
+
+    pub fn str_to_decimal(&self, scale: usize) -> Self {
+        self.inner.clone().str().to_decimal(scale).into()
     }
 
     pub fn str_contains_any(&self, patterns: &RbExpr, ascii_case_insensitive: bool) -> Self {
@@ -333,6 +387,7 @@ impl RbExpr {
         patterns: &RbExpr,
         replace_with: &RbExpr,
         ascii_case_insensitive: bool,
+        leftmost: bool,
     ) -> Self {
         self.inner
             .clone()
@@ -341,7 +396,50 @@ impl RbExpr {
                 patterns.inner.clone(),
                 replace_with.inner.clone(),
                 ascii_case_insensitive,
+                leftmost,
+            )
+            .into()
+    }
+
+    pub fn str_extract_many(
+        &self,
+        patterns: &RbExpr,
+        ascii_case_insensitive: bool,
+        overlapping: bool,
+        leftmost: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .extract_many(
+                patterns.inner.clone(),
+                ascii_case_insensitive,
+                overlapping,
+                leftmost,
+            )
+            .into()
+    }
+
+    pub fn str_find_many(
+        &self,
+        patterns: &RbExpr,
+        ascii_case_insensitive: bool,
+        overlapping: bool,
+        leftmost: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .find_many(
+                patterns.inner.clone(),
+                ascii_case_insensitive,
+                overlapping,
+                leftmost,
             )
             .into()
     }
+
+    pub fn str_escape_regex(&self) -> Self {
+        self.inner.clone().str().escape_regex().into()
+    }
 }

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

@@ -1,4 +1,7 @@
-use crate::RbExpr;
+use magnus::RArray;
+
+use crate::expr::ToExprs;
+use crate::{RbExpr, RbResult};
 
 impl RbExpr {
     pub fn struct_field_by_index(&self, index: i64) -> Self {
@@ -9,6 +12,10 @@ impl RbExpr {
         self.inner.clone().struct_().field_by_name(&name).into()
     }
 
+    pub fn struct_multiple_fields(&self, names: Vec<String>) -> Self {
+        self.inner.clone().struct_().field_by_names(&names).into()
+    }
+
     pub fn struct_rename_fields(&self, names: Vec<String>) -> Self {
         self.inner.clone().struct_().rename_fields(names).into()
     }
@@ -16,4 +23,10 @@ impl RbExpr {
     pub fn struct_json_encode(&self) -> Self {
         self.inner.clone().struct_().json_encode().into()
     }
+
+    pub fn struct_with_fields(&self, fields: RArray) -> RbResult<Self> {
+        let fields = fields.to_exprs()?;
+        let e = self.inner.clone().struct_().with_fields(fields);
+        Ok(e.into())
+    }
 }