polars-df 0.5.0 → 0.6.0

data/ext/polars/src/lib.rs

@@ -17,7 +17,7 @@ mod utils;
 use batched_csv::RbBatchedCsv;
 use conversion::*;
 use dataframe::RbDataFrame;
-use error::{RbPolarsErr, RbValueError};
+use error::{RbPolarsErr, RbTypeError, RbValueError};
 use expr::rb_exprs_to_exprs;
 use expr::RbExpr;
 use functions::whenthen::{RbWhen, RbWhenThen};
@@ -59,7 +59,7 @@ fn init() -> RbResult<()> {
     )?;
     module.define_singleton_method(
         "_concat_lf",
-        function!(crate::functions::lazy::concat_lf, 3),
+        function!(crate::functions::lazy::concat_lf, 4),
     )?;
     module.define_singleton_method(
         "_diag_concat_df",
@@ -179,7 +179,6 @@ fn init() -> RbResult<()> {
         "take_with_series",
         method!(RbDataFrame::take_with_series, 1),
     )?;
-    class.define_method("sort", method!(RbDataFrame::sort, 3))?;
     class.define_method("replace", method!(RbDataFrame::replace, 2))?;
     class.define_method("replace_at_idx", method!(RbDataFrame::replace_at_idx, 2))?;
     class.define_method("insert_at_idx", method!(RbDataFrame::insert_at_idx, 2))?;
@@ -193,7 +192,7 @@ fn init() -> RbResult<()> {
     class.define_method("_clone", method!(RbDataFrame::clone, 0))?;
     class.define_method("melt", method!(RbDataFrame::melt, 4))?;
     class.define_method("pivot_expr", method!(RbDataFrame::pivot_expr, 7))?;
-    class.define_method("partition_by", method!(RbDataFrame::partition_by, 2))?;
+    class.define_method("partition_by", method!(RbDataFrame::partition_by, 3))?;
     class.define_method("shift", method!(RbDataFrame::shift, 1))?;
     class.define_method("lazy", method!(RbDataFrame::lazy, 0))?;
     class.define_method("max", method!(RbDataFrame::max, 0))?;
@@ -208,7 +207,7 @@ fn init() -> RbResult<()> {
     class.define_method("hmin", method!(RbDataFrame::hmin, 0))?;
     class.define_method("hsum", method!(RbDataFrame::hsum, 1))?;
     class.define_method("quantile", method!(RbDataFrame::quantile, 2))?;
-    class.define_method("to_dummies", method!(RbDataFrame::to_dummies, 2))?;
+    class.define_method("to_dummies", method!(RbDataFrame::to_dummies, 3))?;
     class.define_method("null_count", method!(RbDataFrame::null_count, 0))?;
     class.define_method("apply", method!(RbDataFrame::apply, 3))?;
     class.define_method("shrink_to_fit", method!(RbDataFrame::shrink_to_fit, 0))?;
@@ -330,13 +329,14 @@ fn init() -> RbResult<()> {
     class.define_method("cumprod", method!(RbExpr::cumprod, 1))?;
     class.define_method("product", method!(RbExpr::product, 0))?;
     class.define_method("shrink_dtype", method!(RbExpr::shrink_dtype, 0))?;
-    class.define_method("str_parse_date", method!(RbExpr::str_to_date, 4))?;
-    class.define_method("str_parse_datetime", method!(RbExpr::str_to_datetime, 8))?;
-    class.define_method("str_parse_time", method!(RbExpr::str_to_time, 3))?;
+    class.define_method("str_to_date", method!(RbExpr::str_to_date, 4))?;
+    class.define_method("str_to_datetime", method!(RbExpr::str_to_datetime, 6))?;
+    class.define_method("str_to_time", method!(RbExpr::str_to_time, 3))?;
     class.define_method("str_strip", method!(RbExpr::str_strip, 1))?;
     class.define_method("str_rstrip", method!(RbExpr::str_rstrip, 1))?;
     class.define_method("str_lstrip", method!(RbExpr::str_lstrip, 1))?;
     class.define_method("str_slice", method!(RbExpr::str_slice, 2))?;
+    class.define_method("str_explode", method!(RbExpr::str_explode, 0))?;
     class.define_method("str_to_uppercase", method!(RbExpr::str_to_uppercase, 0))?;
     class.define_method("str_to_lowercase", method!(RbExpr::str_to_lowercase, 0))?;
     class.define_method("str_lengths", method!(RbExpr::str_lengths, 0))?;
@@ -349,6 +349,9 @@ fn init() -> RbResult<()> {
     class.define_method("str_contains", method!(RbExpr::str_contains, 3))?;
     class.define_method("str_ends_with", method!(RbExpr::str_ends_with, 1))?;
     class.define_method("str_starts_with", method!(RbExpr::str_starts_with, 1))?;
+    class.define_method("array_max", method!(RbExpr::array_max, 0))?;
+    class.define_method("array_min", method!(RbExpr::array_min, 0))?;
+    class.define_method("array_sum", method!(RbExpr::array_sum, 0))?;
     class.define_method("binary_contains", method!(RbExpr::bin_contains, 1))?;
     class.define_method("binary_ends_with", method!(RbExpr::bin_ends_with, 1))?;
     class.define_method("binary_starts_with", method!(RbExpr::bin_starts_with, 1))?;
@@ -357,7 +360,7 @@ fn init() -> RbResult<()> {
     class.define_method("str_base64_encode", method!(RbExpr::str_base64_encode, 0))?;
     class.define_method("str_base64_decode", method!(RbExpr::str_base64_decode, 1))?;
     class.define_method("str_parse_int", method!(RbExpr::str_parse_int, 2))?;
-    class.define_method("str_json_extract", method!(RbExpr::str_json_extract, 1))?;
+    class.define_method("str_json_extract", method!(RbExpr::str_json_extract, 2))?;
     class.define_method("binary_hex_encode", method!(RbExpr::bin_hex_encode, 0))?;
     class.define_method("binary_hex_decode", method!(RbExpr::bin_hex_decode, 1))?;
     class.define_method(
@@ -457,8 +460,8 @@ fn init() -> RbResult<()> {
     class.define_method("rolling_min", method!(RbExpr::rolling_min, 6))?;
     class.define_method("rolling_max", method!(RbExpr::rolling_max, 6))?;
     class.define_method("rolling_mean", method!(RbExpr::rolling_mean, 6))?;
-    class.define_method("rolling_std", method!(RbExpr::rolling_std, 6))?;
-    class.define_method("rolling_var", method!(RbExpr::rolling_var, 6))?;
+    class.define_method("rolling_std", method!(RbExpr::rolling_std, 7))?;
+    class.define_method("rolling_var", method!(RbExpr::rolling_var, 7))?;
     class.define_method("rolling_median", method!(RbExpr::rolling_median, 6))?;
     class.define_method("rolling_quantile", method!(RbExpr::rolling_quantile, 8))?;
     class.define_method("rolling_skew", method!(RbExpr::rolling_skew, 2))?;
@@ -492,9 +495,9 @@ fn init() -> RbResult<()> {
     class.define_method("reshape", method!(RbExpr::reshape, 1))?;
     class.define_method("cumcount", method!(RbExpr::cumcount, 1))?;
     class.define_method("to_physical", method!(RbExpr::to_physical, 0))?;
-    class.define_method("shuffle", method!(RbExpr::shuffle, 1))?;
-    class.define_method("sample_n", method!(RbExpr::sample_n, 4))?;
-    class.define_method("sample_frac", method!(RbExpr::sample_frac, 4))?;
+    class.define_method("shuffle", method!(RbExpr::shuffle, 2))?;
+    class.define_method("sample_n", method!(RbExpr::sample_n, 5))?;
+    class.define_method("sample_frac", method!(RbExpr::sample_frac, 5))?;
     class.define_method("ewm_mean", method!(RbExpr::ewm_mean, 4))?;
     class.define_method("ewm_std", method!(RbExpr::ewm_std, 5))?;
     class.define_method("ewm_var", method!(RbExpr::ewm_var, 5))?;
@@ -517,6 +520,7 @@ fn init() -> RbResult<()> {
     class.define_method("exp", method!(RbExpr::exp, 0))?;
     class.define_method("entropy", method!(RbExpr::entropy, 2))?;
     class.define_method("_hash", method!(RbExpr::hash, 4))?;
+    class.define_method("set_sorted_flag", method!(RbExpr::set_sorted_flag, 1))?;
 
     // meta
     class.define_method("meta_pop", method!(RbExpr::meta_pop, 0))?;
@@ -541,9 +545,9 @@ fn init() -> RbResult<()> {
     class.define_singleton_method("cols", function!(crate::functions::lazy::cols, 1))?;
     class.define_singleton_method("fold", function!(crate::functions::lazy::fold, 3))?;
     class.define_singleton_method("cumfold", function!(crate::functions::lazy::cumfold, 4))?;
-    class.define_singleton_method("lit", function!(crate::functions::lazy::lit, 1))?;
+    class.define_singleton_method("lit", function!(crate::functions::lazy::lit, 2))?;
     class.define_singleton_method("arange", function!(crate::functions::lazy::arange, 3))?;
-    class.define_singleton_method("repeat", function!(crate::functions::lazy::repeat, 2))?;
+    class.define_singleton_method("repeat", function!(crate::functions::lazy::repeat, 3))?;
     class.define_singleton_method(
         "pearson_corr",
         function!(crate::functions::lazy::pearson_corr, 3),
@@ -589,8 +593,8 @@ fn init() -> RbResult<()> {
         "optimization_toggle",
         method!(RbLazyFrame::optimization_toggle, 7),
     )?;
-    class.define_method("sort", method!(RbLazyFrame::sort, 3))?;
-    class.define_method("sort_by_exprs", method!(RbLazyFrame::sort_by_exprs, 3))?;
+    class.define_method("sort", method!(RbLazyFrame::sort, 4))?;
+    class.define_method("sort_by_exprs", method!(RbLazyFrame::sort_by_exprs, 4))?;
     class.define_method("cache", method!(RbLazyFrame::cache, 0))?;
     class.define_method("collect", method!(RbLazyFrame::collect, 0))?;
     class.define_method("sink_parquet", method!(RbLazyFrame::sink_parquet, 7))?;
@@ -598,7 +602,7 @@ fn init() -> RbResult<()> {
     class.define_method("filter", method!(RbLazyFrame::filter, 1))?;
     class.define_method("select", method!(RbLazyFrame::select, 1))?;
     class.define_method("groupby", method!(RbLazyFrame::groupby, 2))?;
-    class.define_method("groupby_rolling", method!(RbLazyFrame::groupby_rolling, 5))?;
+    class.define_method("groupby_rolling", method!(RbLazyFrame::groupby_rolling, 6))?;
     class.define_method("groupby_dynamic", method!(RbLazyFrame::groupby_dynamic, 9))?;
     class.define_method("with_context", method!(RbLazyFrame::with_context, 1))?;
     class.define_method("join_asof", method!(RbLazyFrame::join_asof, 11))?;
@@ -658,6 +662,7 @@ fn init() -> RbResult<()> {
     class.define_singleton_method("new_null", function!(RbSeries::new_null, 3))?;
     class.define_singleton_method("new_object", function!(RbSeries::new_object, 3))?;
     class.define_singleton_method("new_series_list", function!(RbSeries::new_series_list, 3))?;
+    class.define_singleton_method("new_array", function!(RbSeries::new_array, 5))?;
     class.define_singleton_method("new_decimal", function!(RbSeries::new_decimal, 3))?;
     class.define_singleton_method("repeat", function!(RbSeries::repeat, 4))?;
     class.define_method("struct_unnest", method!(RbSeries::struct_unnest, 0))?;
@@ -725,7 +730,7 @@ fn init() -> RbResult<()> {
     class.define_method("_clone", method!(RbSeries::clone, 0))?;
     class.define_method("apply_lambda", method!(RbSeries::apply_lambda, 3))?;
     class.define_method("zip_with", method!(RbSeries::zip_with, 2))?;
-    class.define_method("to_dummies", method!(RbSeries::to_dummies, 1))?;
+    class.define_method("to_dummies", method!(RbSeries::to_dummies, 2))?;
     class.define_method("peak_max", method!(RbSeries::peak_max, 0))?;
     class.define_method("peak_min", method!(RbSeries::peak_min, 0))?;
     class.define_method("n_unique", method!(RbSeries::n_unique, 0))?;

data/ext/polars/src/object.rs

@@ -10,7 +10,7 @@ use polars_core::prelude::AnyValue;
 use crate::prelude::ObjectValue;
 use crate::Wrap;
 
-// pub(crate) const OBJECT_NAME: &str = "object";
+pub(crate) const OBJECT_NAME: &str = "object";
 
 pub(crate) fn register_object_builder() {
     if !registry::is_object_builder_registered() {

data/ext/polars/src/rb_modules.rs

@@ -12,6 +12,10 @@ pub(crate) fn utils() -> RModule {
     *memoize!(RModule: polars().const_get("Utils").unwrap())
 }
 
+pub(crate) fn bigdecimal() -> RClass {
+    *memoize!(RClass: class::object().const_get("BigDecimal").unwrap())
+}
+
 pub(crate) fn date() -> RClass {
     *memoize!(RClass: class::object().const_get("Date").unwrap())
 }

data/ext/polars/src/series/construction.rs

@@ -1,10 +1,10 @@
 use magnus::RArray;
 use polars_core::prelude::*;
 
-use crate::conversion::{slice_extract_wrapped, Wrap};
+use crate::conversion::{slice_extract_wrapped, vec_extract_wrapped, Wrap};
 use crate::prelude::ObjectValue;
 use crate::series::to_series_collection;
-use crate::{RbPolarsErr, RbResult, RbSeries};
+use crate::{RbPolarsErr, RbResult, RbSeries, RbValueError};
 
 impl RbSeries {
     pub fn new_opt_bool(name: String, obj: RArray, strict: bool) -> RbResult<RbSeries> {
@@ -137,6 +137,32 @@ impl RbSeries {
         Ok(Series::new(&name, &series_vec).into())
     }
 
+    pub fn new_array(
+        width: usize,
+        inner: Option<Wrap<DataType>>,
+        name: String,
+        val: RArray,
+        _strict: bool,
+    ) -> RbResult<Self> {
+        let val = vec_wrap_any_value(val)?;
+        let val = vec_extract_wrapped(val);
+        let out = Series::new(&name, &val);
+        match out.dtype() {
+            DataType::List(list_inner) => {
+                let out = out
+                    .cast(&DataType::Array(
+                        Box::new(inner.map(|dt| dt.0).unwrap_or(*list_inner.clone())),
+                        width,
+                    ))
+                    .map_err(RbPolarsErr::from)?;
+                Ok(out.into())
+            }
+            _ => Err(RbValueError::new_err(
+                "could not create Array from input".to_string(),
+            )),
+        }
+    }
+
     pub fn new_decimal(name: String, val: RArray, strict: bool) -> RbResult<Self> {
         let val = vec_wrap_any_value(val)?;
         // TODO: do we have to respect 'strict' here? it's possible if we want to

data/ext/polars/src/series.rs

@@ -303,7 +303,7 @@ impl RbSeries {
     pub fn to_a(&self) -> Value {
         let series = &self.series.borrow();
 
-        fn to_list_recursive(series: &Series) -> Value {
+        fn to_a_recursive(series: &Series) -> Value {
             let rblist = match series.dtype() {
                 DataType::Boolean => RArray::from_iter(series.bool().unwrap()).into_value(),
                 DataType::UInt8 => RArray::from_iter(series.u8().unwrap()).into_value(),
@@ -339,26 +339,44 @@ impl RbSeries {
                                 v.push(QNIL).unwrap();
                             }
                             Some(s) => {
-                                let rblst = to_list_recursive(s.as_ref());
+                                let rblst = to_a_recursive(s.as_ref());
                                 v.push(rblst).unwrap();
                             }
                         }
                     }
                     v.into_value()
                 }
-                DataType::Date => {
-                    let a = RArray::with_capacity(series.len());
-                    for v in series.iter() {
-                        a.push::<Value>(Wrap(v).into_value()).unwrap();
+                DataType::Array(_, _) => {
+                    let v = RArray::new();
+                    let ca = series.array().unwrap();
+                    for opt_s in ca.amortized_iter() {
+                        match opt_s {
+                            None => {
+                                v.push(QNIL).unwrap();
+                            }
+                            Some(s) => {
+                                let rblst = to_a_recursive(s.as_ref());
+                                v.push(rblst).unwrap();
+                            }
+                        }
                     }
-                    return a.into_value();
+                    v.into_value()
+                }
+                DataType::Date => {
+                    let ca = series.date().unwrap();
+                    return Wrap(ca).into_value();
+                }
+                DataType::Time => {
+                    let ca = series.time().unwrap();
+                    return Wrap(ca).into_value();
                 }
                 DataType::Datetime(_, _) => {
-                    let a = RArray::with_capacity(series.len());
-                    for v in series.iter() {
-                        a.push::<Value>(Wrap(v).into_value()).unwrap();
-                    }
-                    return a.into_value();
+                    let ca = series.datetime().unwrap();
+                    return Wrap(ca).into_value();
+                }
+                DataType::Decimal(_, _) => {
+                    let ca = series.decimal().unwrap();
+                    return Wrap(ca).into_value();
                 }
                 DataType::Utf8 => {
                     let ca = series.utf8().unwrap();
@@ -376,15 +394,37 @@ impl RbSeries {
                     let ca = series.binary().unwrap();
                     return Wrap(ca).into_value();
                 }
-                DataType::Null | DataType::Unknown => {
+                DataType::Null => {
+                    let null: Option<u8> = None;
+                    let n = series.len();
+                    let iter = std::iter::repeat(null).take(n);
+                    use std::iter::{Repeat, Take};
+                    struct NullIter {
+                        iter: Take<Repeat<Option<u8>>>,
+                        n: usize,
+                    }
+                    impl Iterator for NullIter {
+                        type Item = Option<u8>;
+
+                        fn next(&mut self) -> Option<Self::Item> {
+                            self.iter.next()
+                        }
+                        fn size_hint(&self) -> (usize, Option<usize>) {
+                            (self.n, Some(self.n))
+                        }
+                    }
+                    impl ExactSizeIterator for NullIter {}
+
+                    RArray::from_iter(NullIter { iter, n }).into_value()
+                }
+                DataType::Unknown => {
                     panic!("to_a not implemented for null/unknown")
                 }
-                _ => todo!(),
             };
             rblist
         }
 
-        to_list_recursive(series)
+        to_a_recursive(series)
     }
 
     pub fn clone(&self) -> Self {
@@ -594,11 +634,11 @@ impl RbSeries {
         Ok(RbSeries::new(s))
     }
 
-    pub fn to_dummies(&self, sep: Option<String>) -> RbResult<RbDataFrame> {
+    pub fn to_dummies(&self, sep: Option<String>, drop_first: bool) -> RbResult<RbDataFrame> {
         let df = self
             .series
             .borrow()
-            .to_dummies(sep.as_deref())
+            .to_dummies(sep.as_deref(), drop_first)
             .map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }

data/lib/polars/array_expr.rb

@@ -0,0 +1,84 @@
+module Polars
+  # Namespace for array related expressions.
+  class ArrayExpr
+    # @private
+    attr_accessor :_rbexpr
+
+    # @private
+    def initialize(expr)
+      self._rbexpr = expr._rbexpr
+    end
+
+    # Compute the min values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.min)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 1   │
+    #   # │ 3   │
+    #   # └─────┘
+    def min
+      Utils.wrap_expr(_rbexpr.array_min)
+    end
+
+    # Compute the max values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.max)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 2   │
+    #   # │ 4   │
+    #   # └─────┘
+    def max
+      Utils.wrap_expr(_rbexpr.array_max)
+    end
+
+    # Compute the sum values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.sum)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 3   │
+    #   # │ 7   │
+    #   # └─────┘
+    def sum
+      Utils.wrap_expr(_rbexpr.array_sum)
+    end
+  end
+end

data/lib/polars/array_name_space.rb

@@ -0,0 +1,77 @@
+module Polars
+  # Series.arr namespace.
+  class ArrayNameSpace
+    include ExprDispatch
+
+    self._accessor = "arr"
+
+    # @private
+    def initialize(series)
+      self._s = series._s
+    end
+
+    # Compute the min values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     "a", [[1, 2], [4, 3]], dtype: Polars::Array.new(2, Polars::Int64)
+    #   )
+    #   s.arr.min
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [i64]
+    #   # [
+    #   #         1
+    #   #         3
+    #   # ]
+    def min
+      super
+    end
+
+    # Compute the max values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     "a", [[1, 2], [4, 3]], dtype: Polars::Array.new(2, Polars::Int64)
+    #   )
+    #   s.arr.max
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [i64]
+    #   # [
+    #   #         2
+    #   #         4
+    #   # ]
+    def max
+      super
+    end
+
+    # Compute the sum values of the sub-arrays.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.sum)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ i64 │
+    #   # ╞═════╡
+    #   # │ 3   │
+    #   # │ 7   │
+    #   # └─────┘
+    def sum
+      super
+    end
+  end
+end

data/lib/polars/batched_csv_reader.rb

@@ -41,7 +41,7 @@ module Polars
           dtypes.each do|k, v|
             dtype_list << [k, Utils.rb_type_to_dtype(v)]
           end
-        elsif dtypes.is_a?(Array)
+        elsif dtypes.is_a?(::Array)
           dtype_slice = dtypes
         else
           raise ArgumentError, "dtype arg should be list or dict"