polars-df 0.20.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

@@ -6,10 +6,11 @@ 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 {
@@ -655,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()
     }
@@ -868,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;
@@ -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()
@@ -189,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,
@@ -206,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

@@ -57,40 +57,6 @@ 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,

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,11 +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)
+            .to_integer(base.inner.clone(), dtype.map(|wrap| wrap.0), strict)
             .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> {
@@ -42,13 +39,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(())