polars-df 0.1.2 → 0.1.4

data/ext/polars/src/lib.rs

@@ -4,7 +4,10 @@ mod dataframe;
 mod error;
 mod file;
 mod lazy;
+mod list_construction;
 mod series;
+mod set;
+mod utils;
 
 use batched_csv::RbBatchedCsv;
 use conversion::*;
@@ -13,16 +16,32 @@ use error::{RbPolarsErr, RbValueError};
 use file::get_file_like;
 use lazy::dataframe::{RbLazyFrame, RbLazyGroupBy};
 use lazy::dsl::{RbExpr, RbWhen, RbWhenThen};
+use lazy::utils::rb_exprs_to_exprs;
 use magnus::{
     define_module, function, memoize, method, prelude::*, Error, RArray, RClass, RHash, RModule,
     Value,
 };
-use polars::datatypes::DataType;
+use polars::datatypes::{DataType, TimeUnit};
 use polars::error::PolarsResult;
 use polars::frame::DataFrame;
 use polars::functions::{diag_concat_df, hor_concat_df};
+use polars::prelude::{ClosedWindow, Duration, IntoSeries, TimeZone};
 use series::RbSeries;
 
+#[cfg(target_os = "linux")]
+use jemallocator::Jemalloc;
+
+#[cfg(not(target_os = "linux"))]
+use mimalloc::MiMalloc;
+
+#[global_allocator]
+#[cfg(target_os = "linux")]
+static GLOBAL: Jemalloc = Jemalloc;
+
+#[global_allocator]
+#[cfg(not(target_os = "linux"))]
+static GLOBAL: MiMalloc = MiMalloc;
+
 type RbResult<T> = Result<T, Error>;
 
 fn module() -> RModule {
@@ -42,6 +61,9 @@ fn init() -> RbResult<()> {
     module.define_singleton_method("_concat_series", function!(concat_series, 1))?;
     module.define_singleton_method("_ipc_schema", function!(ipc_schema, 1))?;
     module.define_singleton_method("_parquet_schema", function!(parquet_schema, 1))?;
+    module.define_singleton_method("_rb_date_range", function!(rb_date_range, 7))?;
+    module.define_singleton_method("_as_struct", function!(as_struct, 1))?;
+    module.define_singleton_method("_arg_where", function!(arg_where, 1))?;
 
     let class = module.define_class("RbBatchedCsv", Default::default())?;
     class.define_singleton_method("new", function!(RbBatchedCsv::new, -1))?;
@@ -50,7 +72,7 @@ fn init() -> RbResult<()> {
     let class = module.define_class("RbDataFrame", Default::default())?;
     class.define_singleton_method("new", function!(RbDataFrame::init, 1))?;
     class.define_singleton_method("read_csv", function!(RbDataFrame::read_csv, -1))?;
-    class.define_singleton_method("read_parquet", function!(RbDataFrame::read_parquet, 1))?;
+    class.define_singleton_method("read_parquet", function!(RbDataFrame::read_parquet, 7))?;
     class.define_singleton_method("read_ipc", function!(RbDataFrame::read_ipc, 6))?;
     class.define_singleton_method("read_hash", function!(RbDataFrame::read_hash, 1))?;
     class.define_singleton_method("read_json", function!(RbDataFrame::read_json, 1))?;
@@ -60,7 +82,21 @@ fn init() -> RbResult<()> {
     class.define_method("write_ndjson", method!(RbDataFrame::write_ndjson, 1))?;
     class.define_method("write_csv", method!(RbDataFrame::write_csv, 10))?;
     class.define_method("write_ipc", method!(RbDataFrame::write_ipc, 2))?;
+    class.define_method("row_tuple", method!(RbDataFrame::row_tuple, 1))?;
+    class.define_method("row_tuples", method!(RbDataFrame::row_tuples, 0))?;
     class.define_method("write_parquet", method!(RbDataFrame::write_parquet, 5))?;
+    class.define_method("add", method!(RbDataFrame::add, 1))?;
+    class.define_method("sub", method!(RbDataFrame::sub, 1))?;
+    class.define_method("div", method!(RbDataFrame::div, 1))?;
+    class.define_method("mul", method!(RbDataFrame::mul, 1))?;
+    class.define_method("rem", method!(RbDataFrame::rem, 1))?;
+    class.define_method("add_df", method!(RbDataFrame::add_df, 1))?;
+    class.define_method("sub_df", method!(RbDataFrame::sub_df, 1))?;
+    class.define_method("div_df", method!(RbDataFrame::div_df, 1))?;
+    class.define_method("mul_df", method!(RbDataFrame::mul_df, 1))?;
+    class.define_method("rem_df", method!(RbDataFrame::rem_df, 1))?;
+    class.define_method("sample_n", method!(RbDataFrame::sample_n, 4))?;
+    class.define_method("sample_frac", method!(RbDataFrame::sample_frac, 4))?;
     class.define_method("rechunk", method!(RbDataFrame::rechunk, 0))?;
     class.define_method("to_s", method!(RbDataFrame::to_s, 0))?;
     class.define_method("get_columns", method!(RbDataFrame::get_columns, 0))?;
@@ -74,7 +110,19 @@ fn init() -> RbResult<()> {
     class.define_method("shape", method!(RbDataFrame::shape, 0))?;
     class.define_method("height", method!(RbDataFrame::height, 0))?;
     class.define_method("width", method!(RbDataFrame::width, 0))?;
+    class.define_method("hstack_mut", method!(RbDataFrame::hstack_mut, 1))?;
+    class.define_method("hstack", method!(RbDataFrame::hstack, 1))?;
+    class.define_method("extend", method!(RbDataFrame::extend, 1))?;
+    class.define_method("vstack_mut", method!(RbDataFrame::vstack_mut, 1))?;
+    class.define_method("vstack", method!(RbDataFrame::vstack, 1))?;
+    class.define_method("drop_in_place", method!(RbDataFrame::drop_in_place, 1))?;
+    class.define_method("drop_nulls", method!(RbDataFrame::drop_nulls, 1))?;
+    class.define_method("drop", method!(RbDataFrame::drop, 1))?;
     class.define_method("select_at_idx", method!(RbDataFrame::select_at_idx, 1))?;
+    class.define_method(
+        "find_idx_by_name",
+        method!(RbDataFrame::find_idx_by_name, 1),
+    )?;
     class.define_method("column", method!(RbDataFrame::column, 1))?;
     class.define_method("select", method!(RbDataFrame::select, 1))?;
     class.define_method("take", method!(RbDataFrame::take, 1))?;
@@ -116,6 +164,7 @@ fn init() -> RbResult<()> {
     class.define_method("shrink_to_fit", method!(RbDataFrame::shrink_to_fit, 0))?;
     class.define_method("transpose", method!(RbDataFrame::transpose, 2))?;
     class.define_method("upsample", method!(RbDataFrame::upsample, 5))?;
+    class.define_method("to_struct", method!(RbDataFrame::to_struct, 1))?;
     class.define_method("unnest", method!(RbDataFrame::unnest, 1))?;
 
     let class = module.define_class("RbExpr", Default::default())?;
@@ -247,6 +296,14 @@ fn init() -> RbResult<()> {
     class.define_method("str_contains", method!(RbExpr::str_contains, 2))?;
     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("str_hex_encode", method!(RbExpr::str_hex_encode, 0))?;
+    class.define_method("str_hex_decode", method!(RbExpr::str_hex_decode, 1))?;
+    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_json_path_match",
+        method!(RbExpr::str_json_path_match, 1),
+    )?;
     class.define_method("str_extract", method!(RbExpr::str_extract, 2))?;
     class.define_method("str_extract_all", method!(RbExpr::str_extract_all, 1))?;
     class.define_method("count_match", method!(RbExpr::count_match, 1))?;
@@ -304,6 +361,8 @@ fn init() -> RbResult<()> {
     class.define_method("dt_tz_localize", method!(RbExpr::dt_tz_localize, 1))?;
     class.define_method("dt_truncate", method!(RbExpr::dt_truncate, 2))?;
     class.define_method("dt_round", method!(RbExpr::dt_round, 2))?;
+    class.define_method("dot", method!(RbExpr::dot, 1))?;
+    class.define_method("reinterpret", method!(RbExpr::reinterpret, 1))?;
     class.define_method("mode", method!(RbExpr::mode, 0))?;
     class.define_method("keep_name", method!(RbExpr::keep_name, 0))?;
     class.define_method("prefix", method!(RbExpr::prefix, 1))?;
@@ -337,6 +396,7 @@ fn init() -> RbResult<()> {
     class.define_method("lst_slice", method!(RbExpr::lst_slice, 2))?;
     class.define_method("lst_eval", method!(RbExpr::lst_eval, 2))?;
     class.define_method("cumulative_eval", method!(RbExpr::cumulative_eval, 3))?;
+    class.define_method("lst_to_struct", method!(RbExpr::lst_to_struct, 2))?;
     class.define_method("rank", method!(RbExpr::rank, 2))?;
     class.define_method("diff", method!(RbExpr::diff, 2))?;
     class.define_method("pct_change", method!(RbExpr::pct_change, 1))?;
@@ -353,6 +413,7 @@ fn init() -> RbResult<()> {
     class.define_method("ewm_mean", method!(RbExpr::ewm_mean, 3))?;
     class.define_method("ewm_std", method!(RbExpr::ewm_std, 4))?;
     class.define_method("ewm_var", method!(RbExpr::ewm_var, 4))?;
+    class.define_method("extend_constant", method!(RbExpr::extend_constant, 2))?;
     class.define_method("any", method!(RbExpr::any, 0))?;
     class.define_method("all", method!(RbExpr::all, 0))?;
     class.define_method(
@@ -369,6 +430,7 @@ fn init() -> RbResult<()> {
     )?;
     class.define_method("log", method!(RbExpr::log, 1))?;
     class.define_method("exp", method!(RbExpr::exp, 0))?;
+    class.define_method("entropy", method!(RbExpr::entropy, 2))?;
 
     // meta
     class.define_method("meta_pop", method!(RbExpr::meta_pop, 0))?;
@@ -386,6 +448,13 @@ fn init() -> RbResult<()> {
     class.define_singleton_method("fold", function!(crate::lazy::dsl::fold, 3))?;
     class.define_singleton_method("lit", function!(crate::lazy::dsl::lit, 1))?;
     class.define_singleton_method("arange", function!(crate::lazy::dsl::arange, 3))?;
+    class.define_singleton_method("repeat", function!(crate::lazy::dsl::repeat, 2))?;
+    class.define_singleton_method("pearson_corr", function!(crate::lazy::dsl::pearson_corr, 3))?;
+    class.define_singleton_method(
+        "spearman_rank_corr",
+        function!(crate::lazy::dsl::spearman_rank_corr, 4),
+    )?;
+    class.define_singleton_method("cov", function!(crate::lazy::dsl::cov, 2))?;
     class.define_singleton_method("when", function!(crate::lazy::dsl::when, 1))?;
     class.define_singleton_method("concat_str", function!(crate::lazy::dsl::concat_str, 2))?;
     class.define_singleton_method("concat_lst", function!(crate::lazy::dsl::concat_lst, 1))?;
@@ -469,6 +538,10 @@ fn init() -> RbResult<()> {
     class.define_singleton_method("new_opt_f32", function!(RbSeries::new_opt_f32, 3))?;
     class.define_singleton_method("new_opt_f64", function!(RbSeries::new_opt_f64, 3))?;
     class.define_singleton_method("new_str", function!(RbSeries::new_str, 3))?;
+    class.define_singleton_method("new_object", function!(RbSeries::new_object, 3))?;
+    class.define_singleton_method("new_list", function!(RbSeries::new_list, 3))?;
+    class.define_singleton_method("new_opt_date", function!(RbSeries::new_opt_date, 3))?;
+    class.define_singleton_method("new_opt_datetime", function!(RbSeries::new_opt_datetime, 3))?;
     class.define_method("is_sorted_flag", method!(RbSeries::is_sorted_flag, 0))?;
     class.define_method(
         "is_sorted_reverse_flag",
@@ -536,14 +609,85 @@ fn init() -> RbResult<()> {
     class.define_method("kurtosis", method!(RbSeries::kurtosis, 2))?;
     class.define_method("cast", method!(RbSeries::cast, 2))?;
     class.define_method("time_unit", method!(RbSeries::time_unit, 0))?;
-    // rest
-    class.define_method("cumsum", method!(RbSeries::cumsum, 1))?;
-    class.define_method("cummax", method!(RbSeries::cummax, 1))?;
-    class.define_method("cummin", method!(RbSeries::cummin, 1))?;
-    class.define_method("cumprod", method!(RbSeries::cumprod, 1))?;
-    class.define_method("slice", method!(RbSeries::slice, 2))?;
-    class.define_method("ceil", method!(RbSeries::ceil, 0))?;
-    class.define_method("round", method!(RbSeries::round, 1))?;
+    class.define_method("set_at_idx", method!(RbSeries::set_at_idx, 2))?;
+
+    // eq
+    class.define_method("eq_u8", method!(RbSeries::eq_u8, 1))?;
+    class.define_method("eq_u16", method!(RbSeries::eq_u16, 1))?;
+    class.define_method("eq_u32", method!(RbSeries::eq_u32, 1))?;
+    class.define_method("eq_u64", method!(RbSeries::eq_u64, 1))?;
+    class.define_method("eq_i8", method!(RbSeries::eq_i8, 1))?;
+    class.define_method("eq_i16", method!(RbSeries::eq_i16, 1))?;
+    class.define_method("eq_i32", method!(RbSeries::eq_i32, 1))?;
+    class.define_method("eq_i64", method!(RbSeries::eq_i64, 1))?;
+    class.define_method("eq_f32", method!(RbSeries::eq_f32, 1))?;
+    class.define_method("eq_f64", method!(RbSeries::eq_f64, 1))?;
+    // class.define_method("eq_str", method!(RbSeries::eq_str, 1))?;
+
+    // neq
+    class.define_method("neq_u8", method!(RbSeries::neq_u8, 1))?;
+    class.define_method("neq_u16", method!(RbSeries::neq_u16, 1))?;
+    class.define_method("neq_u32", method!(RbSeries::neq_u32, 1))?;
+    class.define_method("neq_u64", method!(RbSeries::neq_u64, 1))?;
+    class.define_method("neq_i8", method!(RbSeries::neq_i8, 1))?;
+    class.define_method("neq_i16", method!(RbSeries::neq_i16, 1))?;
+    class.define_method("neq_i32", method!(RbSeries::neq_i32, 1))?;
+    class.define_method("neq_i64", method!(RbSeries::neq_i64, 1))?;
+    class.define_method("neq_f32", method!(RbSeries::neq_f32, 1))?;
+    class.define_method("neq_f64", method!(RbSeries::neq_f64, 1))?;
+    // class.define_method("neq_str", method!(RbSeries::neq_str, 1))?;
+
+    // gt
+    class.define_method("gt_u8", method!(RbSeries::gt_u8, 1))?;
+    class.define_method("gt_u16", method!(RbSeries::gt_u16, 1))?;
+    class.define_method("gt_u32", method!(RbSeries::gt_u32, 1))?;
+    class.define_method("gt_u64", method!(RbSeries::gt_u64, 1))?;
+    class.define_method("gt_i8", method!(RbSeries::gt_i8, 1))?;
+    class.define_method("gt_i16", method!(RbSeries::gt_i16, 1))?;
+    class.define_method("gt_i32", method!(RbSeries::gt_i32, 1))?;
+    class.define_method("gt_i64", method!(RbSeries::gt_i64, 1))?;
+    class.define_method("gt_f32", method!(RbSeries::gt_f32, 1))?;
+    class.define_method("gt_f64", method!(RbSeries::gt_f64, 1))?;
+    // class.define_method("gt_str", method!(RbSeries::gt_str, 1))?;
+
+    // gt_eq
+    class.define_method("gt_eq_u8", method!(RbSeries::gt_eq_u8, 1))?;
+    class.define_method("gt_eq_u16", method!(RbSeries::gt_eq_u16, 1))?;
+    class.define_method("gt_eq_u32", method!(RbSeries::gt_eq_u32, 1))?;
+    class.define_method("gt_eq_u64", method!(RbSeries::gt_eq_u64, 1))?;
+    class.define_method("gt_eq_i8", method!(RbSeries::gt_eq_i8, 1))?;
+    class.define_method("gt_eq_i16", method!(RbSeries::gt_eq_i16, 1))?;
+    class.define_method("gt_eq_i32", method!(RbSeries::gt_eq_i32, 1))?;
+    class.define_method("gt_eq_i64", method!(RbSeries::gt_eq_i64, 1))?;
+    class.define_method("gt_eq_f32", method!(RbSeries::gt_eq_f32, 1))?;
+    class.define_method("gt_eq_f64", method!(RbSeries::gt_eq_f64, 1))?;
+    // class.define_method("gt_eq_str", method!(RbSeries::gt_eq_str, 1))?;
+
+    // lt
+    class.define_method("lt_u8", method!(RbSeries::lt_u8, 1))?;
+    class.define_method("lt_u16", method!(RbSeries::lt_u16, 1))?;
+    class.define_method("lt_u32", method!(RbSeries::lt_u32, 1))?;
+    class.define_method("lt_u64", method!(RbSeries::lt_u64, 1))?;
+    class.define_method("lt_i8", method!(RbSeries::lt_i8, 1))?;
+    class.define_method("lt_i16", method!(RbSeries::lt_i16, 1))?;
+    class.define_method("lt_i32", method!(RbSeries::lt_i32, 1))?;
+    class.define_method("lt_i64", method!(RbSeries::lt_i64, 1))?;
+    class.define_method("lt_f32", method!(RbSeries::lt_f32, 1))?;
+    class.define_method("lt_f64", method!(RbSeries::lt_f64, 1))?;
+    // class.define_method("lt_str", method!(RbSeries::lt_str, 1))?;
+
+    // lt_eq
+    class.define_method("lt_eq_u8", method!(RbSeries::lt_eq_u8, 1))?;
+    class.define_method("lt_eq_u16", method!(RbSeries::lt_eq_u16, 1))?;
+    class.define_method("lt_eq_u32", method!(RbSeries::lt_eq_u32, 1))?;
+    class.define_method("lt_eq_u64", method!(RbSeries::lt_eq_u64, 1))?;
+    class.define_method("lt_eq_i8", method!(RbSeries::lt_eq_i8, 1))?;
+    class.define_method("lt_eq_i16", method!(RbSeries::lt_eq_i16, 1))?;
+    class.define_method("lt_eq_i32", method!(RbSeries::lt_eq_i32, 1))?;
+    class.define_method("lt_eq_i64", method!(RbSeries::lt_eq_i64, 1))?;
+    class.define_method("lt_eq_f32", method!(RbSeries::lt_eq_f32, 1))?;
+    class.define_method("lt_eq_f64", method!(RbSeries::lt_eq_f64, 1))?;
+    // class.define_method("lt_eq_str", method!(RbSeries::lt_eq_str, 1))?;
 
     let class = module.define_class("RbWhen", Default::default())?;
     class.define_method("_then", method!(RbWhen::then, 1))?;
@@ -641,3 +785,34 @@ fn parquet_schema(rb_f: Value) -> RbResult<Value> {
     }
     Ok(dict.into())
 }
+
+fn rb_date_range(
+    start: i64,
+    stop: i64,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    name: String,
+    tu: Wrap<TimeUnit>,
+    tz: Option<TimeZone>,
+) -> RbSeries {
+    polars::time::date_range_impl(
+        &name,
+        start,
+        stop,
+        Duration::parse(&every),
+        closed.0,
+        tu.0,
+        tz,
+    )
+    .into_series()
+    .into()
+}
+
+fn as_struct(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    Ok(polars::lazy::dsl::as_struct(&exprs).into())
+}
+
+fn arg_where(condition: &RbExpr) -> RbExpr {
+    polars::lazy::dsl::arg_where(condition.inner.clone()).into()
+}

data/ext/polars/src/list_construction.rs

@@ -0,0 +1,100 @@
+use magnus::Value;
+use polars::prelude::*;
+use polars_core::utils::CustomIterTools;
+
+use crate::conversion::get_rbseq;
+use crate::{RbPolarsErr, RbResult};
+
+pub fn rb_seq_to_list(name: &str, seq: Value, dtype: &DataType) -> RbResult<Series> {
+    let (seq, len) = get_rbseq(seq)?;
+
+    let s = match dtype {
+        DataType::Int64 => {
+            let mut builder =
+                ListPrimitiveChunkedBuilder::<Int64Type>::new(name, len, len * 5, DataType::Int64);
+            for sub_seq in seq.each() {
+                let sub_seq = sub_seq?;
+                let (sub_seq, len) = get_rbseq(sub_seq)?;
+
+                // safety: we know the iterators len
+                let iter = unsafe {
+                    sub_seq
+                        .each()
+                        .map(|v| {
+                            let v = v.unwrap();
+                            if v.is_nil() {
+                                None
+                            } else {
+                                Some(v.try_convert::<i64>().unwrap())
+                            }
+                        })
+                        .trust_my_length(len)
+                };
+                builder.append_iter(iter)
+            }
+            builder.finish().into_series()
+        }
+        DataType::Float64 => {
+            let mut builder = ListPrimitiveChunkedBuilder::<Float64Type>::new(
+                name,
+                len,
+                len * 5,
+                DataType::Float64,
+            );
+            for sub_seq in seq.each() {
+                let sub_seq = sub_seq?;
+                let (sub_seq, len) = get_rbseq(sub_seq)?;
+                // safety: we know the iterators len
+                let iter = unsafe {
+                    sub_seq
+                        .each()
+                        .map(|v| {
+                            let v = v.unwrap();
+                            if v.is_nil() {
+                                None
+                            } else {
+                                Some(v.try_convert::<f64>().unwrap())
+                            }
+                        })
+                        .trust_my_length(len)
+                };
+                builder.append_iter(iter)
+            }
+            builder.finish().into_series()
+        }
+        DataType::Boolean => {
+            let mut builder = ListBooleanChunkedBuilder::new(name, len, len * 5);
+            for sub_seq in seq.each() {
+                let sub_seq = sub_seq?;
+                let (sub_seq, len) = get_rbseq(sub_seq)?;
+                // safety: we know the iterators len
+                let iter = unsafe {
+                    sub_seq
+                        .each()
+                        .map(|v| {
+                            let v = v.unwrap();
+                            if v.is_nil() {
+                                None
+                            } else {
+                                Some(v.try_convert::<bool>().unwrap())
+                            }
+                        })
+                        .trust_my_length(len)
+                };
+                builder.append_iter(iter)
+            }
+            builder.finish().into_series()
+        }
+        DataType::Utf8 => {
+            return Err(RbPolarsErr::todo());
+        }
+        dt => {
+            return Err(RbPolarsErr::other(format!(
+                "cannot create list array from {:?}",
+                dt
+            )));
+        }
+    };
+
+    Ok(s)
+}