polars-df 0.13.0 → 0.15.0

data/ext/polars/src/lazyframe/serde.rs

@@ -0,0 +1,31 @@
+use magnus::Value;
+use polars::lazy::frame::LazyFrame;
+use polars::prelude::*;
+use std::io::Read;
+
+use crate::file::get_file_like;
+use crate::{RbLazyFrame, RbResult, RbValueError};
+
+impl RbLazyFrame {
+    // TODO change to serialize_json
+    pub fn read_json(rb_f: Value) -> RbResult<Self> {
+        // it is faster to first read to memory and then parse: https://github.com/serde-rs/json/issues/160
+        // so don't bother with files.
+        let mut json = String::new();
+        let _ = get_file_like(rb_f, false)?
+            .read_to_string(&mut json)
+            .unwrap();
+
+        // Safety
+        // we skipped the serializing/deserializing of the static in lifetime in `DataType`
+        // so we actually don't have a lifetime at all when serializing.
+
+        // &str still has a lifetime. Bit its ok, because we drop it immediately
+        // in this scope
+        let json = unsafe { std::mem::transmute::<&'_ str, &'static str>(json.as_str()) };
+
+        let lp = serde_json::from_str::<DslPlan>(json)
+            .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
+        Ok(LazyFrame::from(lp).into())
+    }
+}

data/ext/polars/src/lib.rs

@@ -3,6 +3,7 @@ mod batched_csv;
 mod conversion;
 mod dataframe;
 mod error;
+mod exceptions;
 mod expr;
 mod file;
 mod functions;
@@ -21,21 +22,31 @@ mod utils;
 use batched_csv::RbBatchedCsv;
 use conversion::*;
 use dataframe::RbDataFrame;
-use error::{RbPolarsErr, RbTypeError, RbValueError};
+use error::RbPolarsErr;
+use exceptions::{RbTypeError, RbValueError};
 use expr::rb_exprs_to_exprs;
 use expr::RbExpr;
 use functions::string_cache::RbStringCacheHolder;
 use functions::whenthen::{RbChainedThen, RbChainedWhen, RbThen, RbWhen};
+use interop::arrow::to_ruby::RbArrowArrayStream;
 use lazyframe::RbLazyFrame;
 use lazygroupby::RbLazyGroupBy;
-use magnus::{define_module, function, method, prelude::*, Error, Ruby};
+use magnus::{define_module, function, method, prelude::*, Ruby};
 use series::RbSeries;
 use sql::RbSQLContext;
 
-type RbResult<T> = Result<T, Error>;
+use magnus::error::Result as RbResult;
+use magnus::Error as RbErr;
+
+// TODO move
+fn re_escape(pattern: String) -> String {
+    regex::escape(&pattern)
+}
 
 #[magnus::init]
 fn init(ruby: &Ruby) -> RbResult<()> {
+    crate::on_startup::register_startup_deps();
+
     let module = define_module("Polars")?;
 
     let class = module.define_class("RbBatchedCsv", ruby.class_object())?;
@@ -45,7 +56,6 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     let class = module.define_class("RbDataFrame", ruby.class_object())?;
     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, 9))?;
     class.define_singleton_method("read_ipc", function!(RbDataFrame::read_ipc, 6))?;
     class.define_singleton_method(
         "read_ipc_stream",
@@ -58,17 +68,21 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_singleton_method("read_ndjson", function!(RbDataFrame::read_ndjson, 4))?;
     class.define_method("estimated_size", method!(RbDataFrame::estimated_size, 0))?;
     class.define_method("dtype_strings", method!(RbDataFrame::dtype_strings, 0))?;
-    class.define_method("write_avro", method!(RbDataFrame::write_avro, 2))?;
+    class.define_method("write_avro", method!(RbDataFrame::write_avro, 3))?;
     class.define_method("write_json", method!(RbDataFrame::write_json, 3))?;
     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("write_ipc", method!(RbDataFrame::write_ipc, 3))?;
     class.define_method(
         "write_ipc_stream",
-        method!(RbDataFrame::write_ipc_stream, 2),
+        method!(RbDataFrame::write_ipc_stream, 3),
     )?;
     class.define_method("row_tuple", method!(RbDataFrame::row_tuple, 1))?;
     class.define_method("row_tuples", method!(RbDataFrame::row_tuples, 0))?;
+    class.define_method(
+        "arrow_c_stream",
+        method!(RbDataFrame::__arrow_c_stream__, 0),
+    )?;
     class.define_method("to_numo", method!(RbDataFrame::to_numo, 0))?;
     class.define_method("write_parquet", method!(RbDataFrame::write_parquet, 6))?;
     class.define_method("add", method!(RbDataFrame::add, 1))?;
@@ -143,6 +157,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("to_struct", method!(RbDataFrame::to_struct, 1))?;
     class.define_method("unnest", method!(RbDataFrame::unnest, 1))?;
     class.define_method("clear", method!(RbDataFrame::clear, 0))?;
+    class.define_method("serialize_json", method!(RbDataFrame::serialize_json, 1))?;
 
     let class = module.define_class("RbExpr", ruby.class_object())?;
     class.define_method("+", method!(RbExpr::add, 1))?;
@@ -286,6 +301,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("str_slice", method!(RbExpr::str_slice, 2))?;
     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_to_titlecase", method!(RbExpr::str_to_titlecase, 0))?;
     class.define_method("str_len_bytes", method!(RbExpr::str_len_bytes, 0))?;
     class.define_method("str_len_chars", method!(RbExpr::str_len_chars, 0))?;
     class.define_method("str_replace_n", method!(RbExpr::str_replace_n, 4))?;
@@ -412,7 +428,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("dt_dst_offset", method!(RbExpr::dt_dst_offset, 0))?;
     class.define_method("dt_round", method!(RbExpr::dt_round, 1))?;
     class.define_method("dt_combine", method!(RbExpr::dt_combine, 2))?;
-    class.define_method("map_batches", method!(RbExpr::map_batches, 4))?;
+    class.define_method("map_batches", method!(RbExpr::map_batches, 5))?;
     class.define_method("dot", method!(RbExpr::dot, 1))?;
     class.define_method("reinterpret", method!(RbExpr::reinterpret, 1))?;
     class.define_method("mode", method!(RbExpr::mode, 0))?;
@@ -567,6 +583,10 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_singleton_method("when", function!(functions::whenthen::when, 1))?;
     class.define_singleton_method("concat_str", function!(functions::lazy::concat_str, 3))?;
     class.define_singleton_method("concat_list", function!(functions::lazy::concat_list, 1))?;
+    class.define_singleton_method(
+        "business_day_count",
+        function!(functions::business::business_day_count, 4),
+    )?;
     class.define_singleton_method(
         "all_horizontal",
         function!(functions::aggregation::all_horizontal, 1),
@@ -699,17 +719,18 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         "set_random_seed",
         function!(functions::random::set_random_seed, 1),
     )?;
+    class.define_singleton_method("re_escape", function!(re_escape, 1))?;
 
     let class = module.define_class("RbLazyFrame", ruby.class_object())?;
     class.define_singleton_method("read_json", function!(RbLazyFrame::read_json, 1))?;
     class.define_singleton_method(
         "new_from_ndjson",
-        function!(RbLazyFrame::new_from_ndjson, 7),
+        function!(RbLazyFrame::new_from_ndjson, 8),
     )?;
     class.define_singleton_method("new_from_csv", function!(RbLazyFrame::new_from_csv, -1))?;
     class.define_singleton_method(
         "new_from_parquet",
-        function!(RbLazyFrame::new_from_parquet, 14),
+        function!(RbLazyFrame::new_from_parquet, -1),
     )?;
     class.define_singleton_method("new_from_ipc", function!(RbLazyFrame::new_from_ipc, 10))?;
     class.define_method("write_json", method!(RbLazyFrame::write_json, 1))?;
@@ -741,15 +762,15 @@ fn init(ruby: &Ruby) -> RbResult<()> {
         method!(RbLazyFrame::group_by_dynamic, 9),
     )?;
     class.define_method("with_context", method!(RbLazyFrame::with_context, 1))?;
-    class.define_method("join_asof", method!(RbLazyFrame::join_asof, 11))?;
-    class.define_method("join", method!(RbLazyFrame::join, 8))?;
+    class.define_method("join_asof", method!(RbLazyFrame::join_asof, 12))?;
+    class.define_method("join", method!(RbLazyFrame::join, 10))?;
     class.define_method("with_column", method!(RbLazyFrame::with_column, 1))?;
     class.define_method("with_columns", method!(RbLazyFrame::with_columns, 1))?;
     class.define_method(
         "with_columns_seq",
         method!(RbLazyFrame::with_columns_seq, 1),
     )?;
-    class.define_method("rename", method!(RbLazyFrame::rename, 2))?;
+    class.define_method("rename", method!(RbLazyFrame::rename, 3))?;
     class.define_method("reverse", method!(RbLazyFrame::reverse, 0))?;
     class.define_method("shift", method!(RbLazyFrame::shift, 2))?;
     class.define_method("fill_nan", method!(RbLazyFrame::fill_nan, 1))?;
@@ -770,6 +791,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("unpivot", method!(RbLazyFrame::unpivot, 4))?;
     class.define_method("with_row_index", method!(RbLazyFrame::with_row_index, 2))?;
     class.define_method("drop", method!(RbLazyFrame::drop, 1))?;
+    class.define_method("cast", method!(RbLazyFrame::cast, 2))?;
     class.define_method("cast_all", method!(RbLazyFrame::cast_all, 2))?;
     class.define_method("_clone", method!(RbLazyFrame::clone, 0))?;
     class.define_method("collect_schema", method!(RbLazyFrame::collect_schema, 0))?;
@@ -810,7 +832,12 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     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_singleton_method(
+        "from_arrow_c_stream",
+        function!(RbSeries::from_arrow_c_stream, 1),
+    )?;
     class.define_method("struct_unnest", method!(RbSeries::struct_unnest, 0))?;
+    class.define_method("struct_fields", method!(RbSeries::struct_fields, 0))?;
     class.define_method(
         "is_sorted_flag",
         method!(RbSeries::is_sorted_ascending_flag, 0),
@@ -1081,5 +1108,9 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     let class = module.define_class("RbStringCacheHolder", ruby.class_object())?;
     class.define_singleton_method("hold", function!(RbStringCacheHolder::hold, 0))?;
 
+    // arrow array stream
+    let class = module.define_class("RbArrowArrayStream", ruby.class_object())?;
+    class.define_method("to_i", method!(RbArrowArrayStream::to_i, 0))?;
+
     Ok(())
 }

data/ext/polars/src/map/dataframe.rs

@@ -7,13 +7,16 @@ use super::*;
 use crate::{RbDataFrame, RbPolarsErr, RbSeries, Wrap};
 
 fn get_iters(df: &DataFrame) -> Vec<SeriesIter> {
-    df.get_columns().iter().map(|s| s.iter()).collect()
+    df.get_columns()
+        .iter()
+        .map(|s| s.as_materialized_series().iter())
+        .collect()
 }
 
 fn get_iters_skip(df: &DataFrame, skip: usize) -> Vec<std::iter::Skip<SeriesIter>> {
     df.get_columns()
         .iter()
-        .map(|s| s.iter().skip(skip))
+        .map(|s| s.as_materialized_series().iter().skip(skip))
         .collect()
 }
 
@@ -113,16 +116,17 @@ pub fn apply_lambda_unknown<'a>(
                 true,
             ));
         } else if out.is_kind_of(class::array()) {
-            return Err(RbPolarsErr::other(
+            return Err(RbPolarsErr::Other(
                 "A list output type is invalid. Do you mean to create polars List Series?\
 Then return a Series object."
                     .into(),
-            ));
+            )
+            .into());
         } else {
-            return Err(RbPolarsErr::other("Could not determine output type".into()));
+            return Err(RbPolarsErr::Other("Could not determine output type".into()).into());
         }
     }
-    Err(RbPolarsErr::other("Could not determine output type".into()))
+    Err(RbPolarsErr::Other("Could not determine output type".into()).into())
 }
 
 fn apply_iter<T>(
@@ -158,10 +162,16 @@ where
 {
     let skip = usize::from(first_value.is_some());
     if init_null_count == df.height() {
-        ChunkedArray::full_null("apply", df.height())
+        ChunkedArray::full_null(PlSmallStr::from_static("map"), df.height())
     } else {
         let iter = apply_iter(df, lambda, init_null_count, skip);
-        iterator_to_primitive(iter, init_null_count, first_value, "apply", df.height())
+        iterator_to_primitive(
+            iter,
+            init_null_count,
+            first_value,
+            PlSmallStr::from_static("map"),
+            df.height(),
+        )
     }
 }
 
@@ -174,10 +184,16 @@ pub fn apply_lambda_with_bool_out_type(
 ) -> ChunkedArray<BooleanType> {
     let skip = usize::from(first_value.is_some());
     if init_null_count == df.height() {
-        ChunkedArray::full_null("apply", df.height())
+        ChunkedArray::full_null(PlSmallStr::from_static("map"), df.height())
     } else {
         let iter = apply_iter(df, lambda, init_null_count, skip);
-        iterator_to_bool(iter, init_null_count, first_value, "apply", df.height())
+        iterator_to_bool(
+            iter,
+            init_null_count,
+            first_value,
+            PlSmallStr::from_static("map"),
+            df.height(),
+        )
     }
 }
 
@@ -190,10 +206,16 @@ pub fn apply_lambda_with_utf8_out_type(
 ) -> StringChunked {
     let skip = usize::from(first_value.is_some());
     if init_null_count == df.height() {
-        ChunkedArray::full_null("apply", df.height())
+        ChunkedArray::full_null(PlSmallStr::from_static("map"), df.height())
     } else {
         let iter = apply_iter::<String>(df, lambda, init_null_count, skip);
-        iterator_to_utf8(iter, init_null_count, first_value, "apply", df.height())
+        iterator_to_utf8(
+            iter,
+            init_null_count,
+            first_value,
+            PlSmallStr::from_static("map"),
+            df.height(),
+        )
     }
 }
 
@@ -207,7 +229,10 @@ pub fn apply_lambda_with_list_out_type(
 ) -> RbResult<ListChunked> {
     let skip = usize::from(first_value.is_some());
     if init_null_count == df.height() {
-        Ok(ChunkedArray::full_null("apply", df.height()))
+        Ok(ChunkedArray::full_null(
+            PlSmallStr::from_static("map"),
+            df.height(),
+        ))
     } else {
         let mut iters = get_iters_skip(df, init_null_count + skip);
         let iter = ((init_null_count + skip)..df.height()).map(|_| {
@@ -229,7 +254,14 @@ pub fn apply_lambda_with_list_out_type(
                 Err(e) => panic!("ruby function failed {}", e),
             }
         });
-        iterator_to_list(dt, iter, init_null_count, first_value, "apply", df.height())
+        iterator_to_list(
+            dt,
+            iter,
+            init_null_count,
+            first_value,
+            PlSmallStr::from_static("map"),
+            df.height(),
+        )
     }
 }
 

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

@@ -1,10 +1,70 @@
-use magnus::Value;
+use magnus::{prelude::*, RArray, Value};
 use polars::prelude::*;
 
-use crate::{RbExpr, Wrap};
+use crate::rb_modules::*;
+use crate::{RbExpr, RbSeries, Wrap};
 
-pub fn binary_lambda(_lambda: Value, _a: Series, _b: Series) -> PolarsResult<Option<Series>> {
-    todo!();
+fn to_series(v: Value, name: &str) -> PolarsResult<Series> {
+    let rb_rbseries = match v.funcall("_s", ()) {
+        Ok(s) => s,
+        // the lambda did not return a series, we try to create a new Ruby Series
+        _ => {
+            let data = RArray::new();
+            data.push(v).unwrap();
+            let res = series().funcall::<_, _, Value>("new", (name.to_string(), data));
+
+            match res {
+                Ok(ruby_s) => ruby_s.funcall::<_, _, &RbSeries>("_s", ()).unwrap(),
+                Err(_) => {
+                    polars_bail!(ComputeError:
+                        "expected a something that could convert to a `Series` but got: {}",
+                        unsafe { v.classname() }
+                    )
+                }
+            }
+        }
+    };
+    // Finally get the actual Series
+    Ok(rb_rbseries.series.borrow().clone())
+}
+
+pub fn binary_lambda(lambda: Value, a: Series, b: Series) -> PolarsResult<Option<Series>> {
+    // create a RbSeries struct/object for Ruby
+    let rbseries_a = RbSeries::new(a);
+    let rbseries_b = RbSeries::new(b);
+
+    // Wrap this RbSeries object in the Ruby side Series wrapper
+    let ruby_series_wrapper_a: Value = utils().funcall("wrap_s", (rbseries_a,)).unwrap();
+    let ruby_series_wrapper_b: Value = utils().funcall("wrap_s", (rbseries_b,)).unwrap();
+
+    // call the lambda and get a Ruby side Series wrapper
+    let result_series_wrapper: Value =
+        match lambda.funcall("call", (ruby_series_wrapper_a, ruby_series_wrapper_b)) {
+            Ok(rbobj) => rbobj,
+            Err(e) => polars_bail!(
+                ComputeError: "custom Ruby function failed: {}", e,
+            ),
+        };
+    let rbseries = if let Ok(rbexpr) = result_series_wrapper.funcall::<_, _, &RbExpr>("_rbexpr", ())
+    {
+        let expr = rbexpr.inner.clone();
+        let df = DataFrame::empty();
+        let out = df
+            .lazy()
+            .select([expr])
+            .with_predicate_pushdown(false)
+            .with_projection_pushdown(false)
+            .collect()?;
+
+        let s = out.select_at_idx(0).unwrap().clone();
+        RbSeries::new(s.take_materialized_series())
+    } else {
+        return Some(to_series(result_series_wrapper, "")).transpose();
+    };
+
+    // Finally get the actual Series
+    let binding = rbseries.series.borrow();
+    Ok(Some(binding.clone()))
 }
 
 pub fn map_single(
@@ -13,6 +73,7 @@ pub fn map_single(
     _output_type: Option<Wrap<DataType>>,
     _agg_list: bool,
     _is_elementwise: bool,
+    _returns_scalar: bool,
 ) -> RbExpr {
     todo!();
 }

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

@@ -28,14 +28,14 @@ fn iterator_to_struct(
     it: impl Iterator<Item = Option<Value>>,
     init_null_count: usize,
     first_value: AnyValue,
-    name: &str,
+    name: PlSmallStr,
     capacity: usize,
 ) -> RbResult<RbSeries> {
     let (vals, flds) = match &first_value {
         av @ AnyValue::Struct(_, _, flds) => (av._iter_struct_av().collect::<Vec<_>>(), &**flds),
         AnyValue::StructOwned(payload) => (payload.0.clone(), &*payload.1),
         _ => {
-            return Err(crate::error::ComputeError::new_err(format!(
+            return Err(crate::exceptions::ComputeError::new_err(format!(
                 "expected struct got {first_value:?}",
             )))
         }
@@ -70,7 +70,7 @@ fn iterator_to_struct(
             Some(dict) => {
                 let dict = RHash::try_convert(dict)?;
                 if dict.len() != struct_width {
-                    return Err(crate::error::ComputeError::new_err(
+                    return Err(crate::exceptions::ComputeError::new_err(
                         format!("Cannot create struct type.\n> The struct dtype expects {} fields, but it got a dict with {} fields.", struct_width, dict.len())
                     ));
                 }
@@ -89,21 +89,23 @@ fn iterator_to_struct(
         items
             .par_iter()
             .zip(flds)
-            .map(|(av, fld)| Series::new(fld.name(), av))
+            .map(|(av, fld)| Series::new(fld.name().clone(), av))
             .collect::<Vec<_>>()
     });
 
-    Ok(StructChunked::from_series(name, &fields)
-        .unwrap()
-        .into_series()
-        .into())
+    Ok(
+        StructChunked::from_series(name, fields[0].len(), fields.iter())
+            .unwrap()
+            .into_series()
+            .into(),
+    )
 }
 
 fn iterator_to_primitive<T>(
     it: impl Iterator<Item = Option<T::Native>>,
     init_null_count: usize,
     first_value: Option<T::Native>,
-    name: &str,
+    name: PlSmallStr,
     capacity: usize,
 ) -> ChunkedArray<T>
 where
@@ -136,7 +138,7 @@ fn iterator_to_bool(
     it: impl Iterator<Item = Option<bool>>,
     init_null_count: usize,
     first_value: Option<bool>,
-    name: &str,
+    name: PlSmallStr,
     capacity: usize,
 ) -> ChunkedArray<BooleanType> {
     // safety: we know the iterators len
@@ -166,7 +168,7 @@ fn iterator_to_object(
     it: impl Iterator<Item = Option<ObjectValue>>,
     init_null_count: usize,
     first_value: Option<ObjectValue>,
-    name: &str,
+    name: PlSmallStr,
     capacity: usize,
 ) -> ObjectChunked<ObjectValue> {
     // safety: we know the iterators len
@@ -196,7 +198,7 @@ fn iterator_to_utf8(
     it: impl Iterator<Item = Option<String>>,
     init_null_count: usize,
     first_value: Option<&str>,
-    name: &str,
+    name: PlSmallStr,
     capacity: usize,
 ) -> StringChunked {
     let first_value = first_value.map(|v| v.to_string());
@@ -229,11 +231,10 @@ fn iterator_to_list(
     it: impl Iterator<Item = Option<Series>>,
     init_null_count: usize,
     first_value: Option<&Series>,
-    name: &str,
+    name: PlSmallStr,
     capacity: usize,
 ) -> RbResult<ListChunked> {
-    let mut builder =
-        get_list_builder(dt, capacity * 5, capacity, name).map_err(RbPolarsErr::from)?;
+    let mut builder = get_list_builder(dt, capacity * 5, capacity, name);
     for _ in 0..init_null_count {
         builder.append_null()
     }
@@ -246,7 +247,7 @@ fn iterator_to_list(
             Some(s) => {
                 if s.len() == 0 && s.dtype() != dt {
                     builder
-                        .append_series(&Series::full_null("", 0, dt))
+                        .append_series(&Series::full_null(PlSmallStr::EMPTY, 0, dt))
                         .unwrap()
                 } else {
                     builder.append_series(&s).map_err(RbPolarsErr::from)?