polars-df 0.5.0 → 0.7.0

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

@@ -1,12 +1,15 @@
-use magnus::{class, RArray, RString, Value};
+use magnus::encoding::{self, EncodingCapable};
+use magnus::{
+    class, prelude::*, typed_data::Obj, value::Opaque, Float, Integer, RArray, RString, Ruby, Value,
+};
 use polars::lazy::dsl;
 use polars::prelude::*;
 
-use crate::apply::lazy::binary_lambda;
 use crate::conversion::{get_lf, get_rbseq, Wrap};
+use crate::map::lazy::binary_lambda;
 use crate::prelude::vec_extract_wrapped;
 use crate::rb_exprs_to_exprs;
-use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
+use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
 macro_rules! set_unwrapped_or_0 {
     ($($var:ident),+ $(,)?) => {
@@ -14,10 +17,6 @@ macro_rules! set_unwrapped_or_0 {
     };
 }
 
-pub fn arange(low: &RbExpr, high: &RbExpr, step: i64) -> RbExpr {
-    dsl::arange(low.inner.clone(), high.inner.clone(), step).into()
-}
-
 pub fn arg_sort_by(by: RArray, descending: Vec<bool>) -> RbResult<RbExpr> {
     let by = rb_exprs_to_exprs(by)?;
     Ok(dsl::arg_sort_by(by, &descending).into())
@@ -29,7 +28,7 @@ pub fn arg_where(condition: &RbExpr) -> RbExpr {
 
 pub fn as_struct(exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
-    Ok(dsl::as_struct(&exprs).into())
+    Ok(dsl::as_struct(exprs).into())
 }
 
 pub fn coalesce(exprs: RArray) -> RbResult<RbExpr> {
@@ -44,7 +43,7 @@ pub fn col(name: String) -> RbExpr {
 pub fn collect_all(lfs: RArray) -> RbResult<RArray> {
     let lfs = lfs
         .each()
-        .map(|v| v?.try_convert::<&RbLazyFrame>())
+        .map(|v| <&RbLazyFrame>::try_convert(v?))
         .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
 
     Ok(RArray::from_iter(lfs.iter().map(|lf| {
@@ -57,7 +56,12 @@ pub fn cols(names: Vec<String>) -> RbExpr {
     dsl::cols(names).into()
 }
 
-pub fn concat_lf(lfs: Value, rechunk: bool, parallel: bool) -> RbResult<RbLazyFrame> {
+pub fn concat_lf(
+    lfs: Value,
+    rechunk: bool,
+    parallel: bool,
+    to_supertypes: bool,
+) -> RbResult<RbLazyFrame> {
     let (seq, len) = get_rbseq(lfs)?;
     let mut lfs = Vec::with_capacity(len);
 
@@ -67,40 +71,77 @@ pub fn concat_lf(lfs: Value, rechunk: bool, parallel: bool) -> RbResult<RbLazyFr
         lfs.push(lf);
     }
 
-    let lf = polars::lazy::dsl::concat(lfs, rechunk, parallel).map_err(RbPolarsErr::from)?;
+    let lf = dsl::concat(
+        lfs,
+        UnionArgs {
+            rechunk,
+            parallel,
+            to_supertypes,
+        },
+    )
+    .map_err(RbPolarsErr::from)?;
+    Ok(lf.into())
+}
+
+pub fn concat_lf_diagonal(
+    lfs: RArray,
+    rechunk: bool,
+    parallel: bool,
+    to_supertypes: bool,
+) -> RbResult<RbLazyFrame> {
+    let iter = lfs.each();
+
+    let lfs = iter
+        .map(|item| {
+            let item = item?;
+            get_lf(item)
+        })
+        .collect::<RbResult<Vec<_>>>()?;
+
+    let lf = dsl::functions::concat_lf_diagonal(
+        lfs,
+        UnionArgs {
+            rechunk,
+            parallel,
+            to_supertypes,
+        },
+    )
+    .map_err(RbPolarsErr::from)?;
     Ok(lf.into())
 }
 
 #[allow(clippy::too_many_arguments)]
 pub fn duration(
+    weeks: Option<&RbExpr>,
     days: Option<&RbExpr>,
+    hours: Option<&RbExpr>,
+    minutes: Option<&RbExpr>,
     seconds: Option<&RbExpr>,
-    nanoseconds: Option<&RbExpr>,
-    microseconds: Option<&RbExpr>,
     milliseconds: Option<&RbExpr>,
-    minutes: Option<&RbExpr>,
-    hours: Option<&RbExpr>,
-    weeks: Option<&RbExpr>,
+    microseconds: Option<&RbExpr>,
+    nanoseconds: Option<&RbExpr>,
+    time_unit: Wrap<TimeUnit>,
 ) -> RbExpr {
     set_unwrapped_or_0!(
+        weeks,
         days,
+        hours,
+        minutes,
         seconds,
-        nanoseconds,
-        microseconds,
         milliseconds,
-        minutes,
-        hours,
-        weeks,
+        microseconds,
+        nanoseconds,
     );
     let args = DurationArgs {
+        weeks,
         days,
+        hours,
+        minutes,
         seconds,
-        nanoseconds,
-        microseconds,
         milliseconds,
-        minutes,
-        hours,
-        weeks,
+        microseconds,
+        nanoseconds,
+        time_unit: time_unit.0,
     };
     dsl::duration(args).into()
 }
@@ -123,28 +164,27 @@ pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
 
 pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
+    let lambda = Opaque::from(lambda);
 
-    let func = move |a: Series, b: Series| binary_lambda(lambda, a, b);
+    let func =
+        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
     Ok(polars::lazy::dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
 }
 
 pub fn cumfold(acc: &RbExpr, lambda: Value, exprs: RArray, include_init: bool) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
+    let lambda = Opaque::from(lambda);
 
-    let func = move |a: Series, b: Series| binary_lambda(lambda, a, b);
-    Ok(polars::lazy::dsl::cumfold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+    let func =
+        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
+    Ok(polars::lazy::dsl::cum_fold_exprs(acc.inner.clone(), func, exprs, include_init).into())
 }
 
-// TODO improve
-pub fn lit(value: Value) -> RbResult<RbExpr> {
-    if value.is_nil() {
-        Ok(dsl::lit(Null {}).into())
-    } else if let Ok(series) = value.try_convert::<&RbSeries>() {
-        Ok(dsl::lit(series.series.borrow().clone()).into())
-    } else if let Some(v) = RString::from_value(value) {
-        Ok(dsl::lit(v.try_convert::<String>()?).into())
-    } else if value.is_kind_of(class::integer()) {
-        match value.try_convert::<i64>() {
+pub fn lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
+    if value.is_kind_of(class::true_class()) || value.is_kind_of(class::false_class()) {
+        Ok(dsl::lit(bool::try_convert(value)?).into())
+    } else if let Some(v) = Integer::from_value(value) {
+        match v.to_i64() {
             Ok(val) => {
                 if val > 0 && val < i32::MAX as i64 || val < 0 && val > i32::MIN as i64 {
                     Ok(dsl::lit(val as i32).into())
@@ -153,21 +193,51 @@ pub fn lit(value: Value) -> RbResult<RbExpr> {
                 }
             }
             _ => {
-                let val = value.try_convert::<u64>()?;
+                let val = v.to_u64()?;
                 Ok(dsl::lit(val).into())
             }
         }
+    } else if let Some(v) = Float::from_value(value) {
+        Ok(dsl::lit(v.to_f64()).into())
+    } else if let Some(v) = RString::from_value(value) {
+        if v.enc_get() == encoding::Index::utf8() {
+            Ok(dsl::lit(v.to_string()?).into())
+        } else {
+            Ok(dsl::lit(unsafe { v.as_slice() }).into())
+        }
+    } else if let Ok(series) = Obj::<RbSeries>::try_convert(value) {
+        Ok(dsl::lit(series.series.borrow().clone()).into())
+    } else if value.is_nil() {
+        Ok(dsl::lit(Null {}).into())
+    } else if allow_object {
+        todo!()
     } else {
-        Ok(dsl::lit(value.try_convert::<f64>()?).into())
+        Err(RbValueError::new_err(format!(
+            "could not convert value {:?} as a Literal",
+            value.to_string()
+        )))
     }
 }
 
-pub fn repeat(value: Value, n_times: &RbExpr) -> RbResult<RbExpr> {
-    if value.is_nil() {
-        Ok(polars::lazy::dsl::repeat(Null {}, n_times.inner.clone()).into())
-    } else {
-        todo!();
+pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbResult<RbExpr> {
+    let mut value = value.inner.clone();
+    let n = n.inner.clone();
+
+    if let Some(dtype) = dtype {
+        value = value.cast(dtype.0);
+    }
+
+    if let Expr::Literal(lv) = &value {
+        let av = lv.to_anyvalue().unwrap();
+        // Integer inputs that fit in Int32 are parsed as such
+        if let DataType::Int64 = av.dtype() {
+            let int_value = av.try_extract::<i64>().unwrap();
+            if int_value >= i32::MIN as i64 && int_value <= i32::MAX as i64 {
+                value = value.cast(DataType::Int32);
+            }
+        }
     }
+    Ok(dsl::repeat(value, n).into())
 }
 
 pub fn pearson_corr(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
@@ -179,8 +249,8 @@ pub fn spearman_rank_corr(a: &RbExpr, b: &RbExpr, ddof: u8, propagate_nans: bool
         .into()
 }
 
-pub fn cov(a: &RbExpr, b: &RbExpr) -> RbExpr {
-    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone()).into()
+pub fn cov(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
+    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone(), ddof).into()
 }
 
 pub fn concat_str(s: RArray, sep: String) -> RbResult<RbExpr> {
@@ -197,13 +267,8 @@ pub fn concat_lst(s: RArray) -> RbResult<RbExpr> {
 pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
     let dtypes = dtypes
         .each()
-        .map(|v| v?.try_convert::<Wrap<DataType>>())
+        .map(|v| Wrap::<DataType>::try_convert(v?))
         .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
     let dtypes = vec_extract_wrapped(dtypes);
     Ok(crate::functions::lazy::dtype_cols(dtypes))
 }
-
-pub fn sum_exprs(exprs: RArray) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
-    Ok(polars::lazy::dsl::sum_exprs(exprs).into())
-}

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

@@ -1,8 +1,38 @@
 use magnus::{IntoValue, Value};
+use polars_core;
+use polars_core::fmt::FloatFmt;
 use polars_core::prelude::IDX_DTYPE;
+use polars_core::POOL;
 
 use crate::conversion::Wrap;
+use crate::{RbResult, RbValueError};
 
 pub fn get_idx_type() -> Value {
     Wrap(IDX_DTYPE).into_value()
 }
+
+pub fn threadpool_size() -> usize {
+    POOL.current_num_threads()
+}
+
+pub fn set_float_fmt(fmt: String) -> RbResult<()> {
+    let fmt = match fmt.as_str() {
+        "full" => FloatFmt::Full,
+        "mixed" => FloatFmt::Mixed,
+        e => {
+            return Err(RbValueError::new_err(format!(
+                "fmt must be one of {{'full', 'mixed'}}, got {e}",
+            )))
+        }
+    };
+    polars_core::fmt::set_float_fmt(fmt);
+    Ok(())
+}
+
+pub fn get_float_fmt() -> RbResult<String> {
+    let strfmt = match polars_core::fmt::get_float_fmt() {
+        FloatFmt::Full => "full",
+        FloatFmt::Mixed => "mixed",
+    };
+    Ok(strfmt.to_string())
+}

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

@@ -0,0 +1,8 @@
+use crate::conversion::Wrap;
+use crate::prelude::DataType;
+use crate::RbResult;
+
+pub fn dtype_str_repr(dtype: Wrap<DataType>) -> RbResult<String> {
+    let dtype = dtype.0;
+    Ok(dtype.to_string())
+}

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

@@ -1,5 +1,10 @@
+pub mod aggregation;
 pub mod eager;
 pub mod io;
 pub mod lazy;
 pub mod meta;
+pub mod misc;
+pub mod random;
+pub mod range;
+pub mod string_cache;
 pub mod whenthen;

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

@@ -0,0 +1,6 @@
+use crate::RbResult;
+
+pub fn set_random_seed(seed: u64) -> RbResult<()> {
+    polars_core::random::set_global_random_seed(seed);
+    Ok(())
+}

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

@@ -0,0 +1,46 @@
+use polars::lazy::dsl;
+use polars_core::datatypes::{TimeUnit, TimeZone};
+
+use crate::conversion::Wrap;
+use crate::prelude::*;
+use crate::RbExpr;
+
+pub fn int_range(start: &RbExpr, end: &RbExpr, step: i64, dtype: Wrap<DataType>) -> RbExpr {
+    let dtype = dtype.0;
+
+    let mut result = dsl::int_range(start.inner.clone(), end.inner.clone(), step);
+
+    if dtype != DataType::Int64 {
+        result = result.cast(dtype)
+    }
+
+    result.into()
+}
+
+pub fn int_ranges(start: &RbExpr, end: &RbExpr, step: i64, dtype: Wrap<DataType>) -> RbExpr {
+    let dtype = dtype.0;
+
+    let mut result = dsl::int_ranges(start.inner.clone(), end.inner.clone(), step);
+
+    if dtype != DataType::Int64 {
+        result = result.cast(DataType::List(Box::new(dtype)))
+    }
+
+    result.into()
+}
+
+pub fn date_range(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    time_unit: Option<Wrap<TimeUnit>>,
+    time_zone: Option<TimeZone>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    let time_unit = time_unit.map(|x| x.0);
+    dsl::date_range(start, end, every, closed, time_unit, time_zone).into()
+}

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

@@ -0,0 +1,11 @@
+pub fn enable_string_cache() {
+    polars_core::enable_string_cache()
+}
+
+pub fn disable_string_cache() {
+    polars_core::disable_string_cache()
+}
+
+pub fn using_string_cache() -> bool {
+    polars_core::using_string_cache()
+}

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

@@ -16,23 +16,23 @@ impl From<dsl::When> for RbWhen {
 
 #[magnus::wrap(class = "Polars::RbWhenThen")]
 #[derive(Clone)]
-pub struct RbWhenThen {
-    pub inner: dsl::WhenThen,
+pub struct RbThen {
+    pub inner: dsl::Then,
 }
 
-impl From<dsl::WhenThen> for RbWhenThen {
-    fn from(inner: dsl::WhenThen) -> Self {
-        RbWhenThen { inner }
+impl From<dsl::Then> for RbThen {
+    fn from(inner: dsl::Then) -> Self {
+        RbThen { inner }
     }
 }
 
 impl RbWhen {
-    pub fn then(&self, expr: &RbExpr) -> RbWhenThen {
+    pub fn then(&self, expr: &RbExpr) -> RbThen {
         self.inner.clone().then(expr.inner.clone()).into()
     }
 }
 
-impl RbWhenThen {
+impl RbThen {
     pub fn overwise(&self, expr: &RbExpr) -> RbExpr {
         self.inner.clone().otherwise(expr.inner.clone()).into()
     }