polars-df 0.20.0 → 0.21.1

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;
@@ -47,6 +47,14 @@ impl RbExpr {
         self.inner.clone().list().eval(expr.inner.clone()).into()
     }
 
+    pub fn list_filter(&self, predicate: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .eval(Expr::Column(PlSmallStr::EMPTY).filter(predicate.inner.clone()))
+            .into()
+    }
+
     pub fn list_get(&self, index: &RbExpr, null_on_oob: bool) -> Self {
         self.inner
             .clone()
@@ -75,6 +83,18 @@ impl RbExpr {
         self.inner.clone().list().mean().into()
     }
 
+    pub fn list_median(&self) -> Self {
+        self.inner.clone().list().median().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()
     }
@@ -163,6 +183,14 @@ 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()
     }
@@ -189,7 +217,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,
@@ -197,6 +225,10 @@ impl RbExpr {
             .into())
     }
 
+    pub fn list_n_unique(&self) -> Self {
+        self.inner.clone().list().n_unique().into()
+    }
+
     pub fn list_unique(&self, maintain_order: bool) -> Self {
         let e = self.inner.clone();
 
@@ -206,4 +238,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

@@ -57,38 +57,15 @@ 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
+    pub fn meta_is_column_selection(&self, allow_aliasing: bool) -> bool {
+        self.inner
             .clone()
             .meta()
-            ._selector_and(other.inner.clone())
-            .map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+            .is_column_selection(allow_aliasing)
     }
 
-    pub fn _meta_as_selector(&self) -> RbExpr {
-        self.inner.clone().meta()._into_selector().into()
+    pub fn meta_is_literal(&self, allow_aliasing: bool) -> bool {
+        self.inner.clone().meta().is_literal(allow_aliasing)
     }
 
     fn compute_tree_format(

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

@@ -1,16 +1,19 @@
 mod array;
 mod binary;
+mod bitwise;
 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(),
                     )),
                 }
             })
@@ -42,4 +42,12 @@ impl RbExpr {
     pub fn name_to_uppercase(&self) -> Self {
         self.inner.clone().name().to_uppercase().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,7 @@
 use polars::prelude::*;
 
-use crate::conversion::Wrap;
 use crate::RbExpr;
+use crate::conversion::Wrap;
 
 impl RbExpr {
     pub fn rolling_sum(
@@ -337,4 +337,24 @@ impl RbExpr {
 
         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()
+    }
 }

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

@@ -122,6 +122,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 +167,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 +212,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,11 +257,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)
+            .to_integer(base.inner.clone(), dtype.map(|wrap| wrap.0), strict)
             .into()
     }
 
@@ -343,4 +380,34 @@ impl RbExpr {
             )
             .into()
     }
+
+    pub fn str_extract_many(
+        &self,
+        patterns: &RbExpr,
+        ascii_case_insensitive: bool,
+        overlapping: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .extract_many(patterns.inner.clone(), ascii_case_insensitive, overlapping)
+            .into()
+    }
+
+    pub fn str_find_many(
+        &self,
+        patterns: &RbExpr,
+        ascii_case_insensitive: bool,
+        overlapping: bool,
+    ) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .find_many(patterns.inner.clone(), ascii_case_insensitive, overlapping)
+            .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,6 @@
-use crate::RbExpr;
+use magnus::RArray;
+
+use crate::{RbExpr, RbResult, rb_exprs_to_exprs};
 
 impl RbExpr {
     pub fn struct_field_by_index(&self, index: i64) -> Self {
@@ -16,4 +18,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 = rb_exprs_to_exprs(fields)?;
+        let e = self.inner.clone().struct_().with_fields(fields);
+        Ok(e.into())
+    }
 }

data/ext/polars/src/file.rs

@@ -3,17 +3,18 @@ use std::io;
 use std::io::{Cursor, Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 
-use magnus::{exception, prelude::*, value::Opaque, Error, RString, Ruby, Value};
+use magnus::{Error, RString, Ruby, Value, exception, prelude::*, value::Opaque};
 use polars::io::cloud::CloudOptions;
 use polars::io::mmap::MmapBytesReader;
+use polars::prelude::PlPath;
 use polars::prelude::file::DynWriteable;
 use polars::prelude::sync_on_close::SyncOnCloseType;
 use polars_utils::file::ClosableFile;
 use polars_utils::mmap::MemSlice;
 
+use crate::RbResult;
 use crate::error::RbPolarsErr;
 use crate::prelude::resolve_homedir;
-use crate::RbResult;
 
 #[derive(Clone)]
 pub struct RbFileLikeObject {
@@ -188,7 +189,7 @@ impl EitherRustRubyFile {
 
 pub enum RubyScanSourceInput {
     Buffer(MemSlice),
-    Path(PathBuf),
+    Path(PlPath),
     #[allow(dead_code)]
     File(File),
 }
@@ -202,8 +203,13 @@ pub(crate) fn try_get_rbfile(
 }
 
 pub fn get_ruby_scan_source_input(rb_f: Value, write: bool) -> RbResult<RubyScanSourceInput> {
-    if let Ok(file_path) = PathBuf::try_convert(rb_f) {
-        // TODO resolve_homedir
+    if let Ok(s) = String::try_convert(rb_f) {
+        let mut file_path = PlPath::new(&s);
+        if let Some(p) = file_path.as_ref().as_local_path() {
+            if p.starts_with("~/") {
+                file_path = PlPath::Local(resolve_homedir(&p).into());
+            }
+        }
         Ok(RubyScanSourceInput::Path(file_path))
     } else {
         let f = RbFileLikeObject::with_requirements(rb_f, !write, write, !write)?;

data/ext/polars/src/functions/io.rs

@@ -1,13 +1,10 @@
 use std::io::BufReader;
 
-use arrow::array::Utf8ViewArray;
 use magnus::{RHash, Value};
 use polars::prelude::ArrowSchema;
-use polars_core::datatypes::create_enum_dtype;
 
 use crate::conversion::Wrap;
-use crate::file::{get_either_file, EitherRustRubyFile};
-use crate::prelude::ArrowDataType;
+use crate::file::{EitherRustRubyFile, get_either_file};
 use crate::{RbPolarsErr, RbResult};
 
 pub fn read_ipc_schema(rb_f: Value) -> RbResult<RHash> {
@@ -24,6 +21,25 @@ pub fn read_ipc_schema(rb_f: Value) -> RbResult<RHash> {
     Ok(dict)
 }
 
+pub fn read_parquet_metadata(rb_f: Value) -> RbResult<RHash> {
+    use polars_parquet::read::read_metadata;
+    use polars_parquet::read::schema::read_custom_key_value_metadata;
+
+    let metadata = match get_either_file(rb_f, false)? {
+        EitherRustRubyFile::Rust(r) => {
+            read_metadata(&mut BufReader::new(r)).map_err(RbPolarsErr::from)?
+        }
+        EitherRustRubyFile::Rb(mut r) => read_metadata(&mut r).map_err(RbPolarsErr::from)?,
+    };
+
+    let key_value_metadata = read_custom_key_value_metadata(metadata.key_value_metadata());
+    let dict = RHash::new();
+    for (key, value) in key_value_metadata.into_iter() {
+        dict.aset(key.as_str(), value.as_str())?;
+    }
+    Ok(dict)
+}
+
 pub fn read_parquet_schema(rb_f: Value) -> RbResult<RHash> {
     use polars_parquet::read::{infer_schema, read_metadata};
 
@@ -42,13 +58,7 @@ pub fn read_parquet_schema(rb_f: Value) -> RbResult<RHash> {
 
 fn fields_to_rbdict(schema: &ArrowSchema, dict: &RHash) -> RbResult<()> {
     for field in schema.iter_values() {
-        let dt = if field.is_enum() {
-            Wrap(create_enum_dtype(Utf8ViewArray::new_empty(
-                ArrowDataType::Utf8View,
-            )))
-        } else {
-            Wrap(polars::prelude::DataType::from_arrow_field(field))
-        };
+        let dt = Wrap(polars::prelude::DataType::from_arrow_field(field));
         dict.aset(field.name.as_str(), dt)?;
     }
     Ok(())