polars-df 0.15.0 → 0.26.0

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

@@ -1,79 +1,49 @@
-use magnus::{prelude::*, RArray, Value};
+use magnus::{KwArgs, RArray, Ruby, Value, prelude::*, value::Opaque};
 use polars::prelude::*;
 
-use crate::rb_modules::*;
-use crate::{RbExpr, RbSeries, Wrap};
-
-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()))
+use crate::expr::ToExprs;
+use crate::expr::datatype::RbDataTypeExpr;
+use crate::ruby::ruby_udf::{RubyUdfExpression, RubyUdfExt};
+use crate::ruby::utils::{TryIntoValue, to_pl_err};
+use crate::{RbExpr, RbResult, RbSeries, Wrap};
+
+pub(crate) fn call_lambda_with_series(
+    rb: &Ruby,
+    s: &[Column],
+    output_dtype: Option<DataType>,
+    lambda: Opaque<Value>,
+) -> PolarsResult<Column> {
+    let lambda = rb.get_inner(lambda);
+
+    // Set return_dtype in kwargs
+    let dict = rb.hash_new();
+    let output_dtype = output_dtype
+        .map(|v| Wrap(v).try_into_value_with(rb))
+        .transpose()
+        .map_err(to_pl_err)?;
+    dict.aset(rb.sym_new("return_dtype"), output_dtype)
+        .map_err(to_pl_err)?;
+
+    let series_objects = rb.ary_from_iter(
+        s.iter()
+            .map(|c| RbSeries::new(c.as_materialized_series().clone())),
+    );
+
+    let result = lambda.funcall::<_, _, &RbSeries>("call", (series_objects, KwArgs(dict)));
+    result
+        .map_err(to_pl_err)
+        .map(|s| s.clone().series.into_inner().into_column())
 }
 
-pub fn map_single(
-    _rbexpr: &RbExpr,
-    _lambda: Value,
-    _output_type: Option<Wrap<DataType>>,
-    _agg_list: bool,
-    _is_elementwise: bool,
-    _returns_scalar: bool,
-) -> RbExpr {
-    todo!();
+pub fn map_expr(
+    rbexpr: RArray,
+    lambda: Value,
+    output_type: Option<&RbDataTypeExpr>,
+    is_elementwise: bool,
+    returns_scalar: bool,
+) -> RbResult<RbExpr> {
+    let output_type = output_type.map(|v| v.inner.clone());
+    let func = RubyUdfExpression::new(lambda, output_type, is_elementwise, returns_scalar);
+    let exprs = rbexpr.to_exprs()?;
+    Ok(Expr::map_many_ruby(exprs, func).into())
 }

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

@@ -1,259 +1,23 @@
-pub mod dataframe;
 pub mod lazy;
 pub mod series;
 
-use magnus::{prelude::*, RHash, Value};
-use polars::chunked_array::builder::get_list_builder;
+use magnus::{Ruby, Value, prelude::*};
 use polars::prelude::*;
-use polars_core::export::rayon::prelude::*;
-use polars_core::utils::CustomIterTools;
-use polars_core::POOL;
 
-use crate::{ObjectValue, RbPolarsErr, RbResult, RbSeries, Wrap};
-
-pub trait RbArrowPrimitiveType: PolarsNumericType {}
-
-impl RbArrowPrimitiveType for UInt8Type {}
-impl RbArrowPrimitiveType for UInt16Type {}
-impl RbArrowPrimitiveType for UInt32Type {}
-impl RbArrowPrimitiveType for UInt64Type {}
-impl RbArrowPrimitiveType for Int8Type {}
-impl RbArrowPrimitiveType for Int16Type {}
-impl RbArrowPrimitiveType for Int32Type {}
-impl RbArrowPrimitiveType for Int64Type {}
-impl RbArrowPrimitiveType for Float32Type {}
-impl RbArrowPrimitiveType for Float64Type {}
-
-fn iterator_to_struct(
-    it: impl Iterator<Item = Option<Value>>,
-    init_null_count: usize,
-    first_value: AnyValue,
-    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::exceptions::ComputeError::new_err(format!(
-                "expected struct got {first_value:?}",
-            )))
-        }
-    };
-
-    let struct_width = vals.len();
-
-    // every item in the struct is kept as its own buffer of anyvalues
-    // so as struct with 2 items: {a, b}
-    // will have
-    // [
-    //      [ a values ]
-    //      [ b values ]
-    // ]
-    let mut items = Vec::with_capacity(vals.len());
-    for item in vals {
-        let mut buf = Vec::with_capacity(capacity);
-        for _ in 0..init_null_count {
-            buf.push(AnyValue::Null);
-        }
-        buf.push(item.clone());
-        items.push(buf);
-    }
-
-    for dict in it {
-        match dict {
-            None => {
-                for field_items in &mut items {
-                    field_items.push(AnyValue::Null);
-                }
-            }
-            Some(dict) => {
-                let dict = RHash::try_convert(dict)?;
-                if dict.len() != struct_width {
-                    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())
-                    ));
-                }
-                // we ignore the keys of the rest of the dicts
-                // the first item determines the output name
-                todo!()
-                // for ((_, val), field_items) in dict.iter().zip(&mut items) {
-                //     let item = Wrap::<AnyValue>::try_convert(val)?;
-                //     field_items.push(item.0)
-                // }
-            }
-        }
-    }
-
-    let fields = POOL.install(|| {
-        items
-            .par_iter()
-            .zip(flds)
-            .map(|(av, fld)| Series::new(fld.name().clone(), av))
-            .collect::<Vec<_>>()
-    });
-
-    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: PlSmallStr,
-    capacity: usize,
-) -> ChunkedArray<T>
-where
-    T: RbArrowPrimitiveType,
-{
-    // safety: we know the iterators len
-    let mut ca: ChunkedArray<T> = unsafe {
-        if init_null_count > 0 {
-            (0..init_null_count)
-                .map(|_| None)
-                .chain(std::iter::once(first_value))
-                .chain(it)
-                .trust_my_length(capacity)
-                .collect_trusted()
-        } else if first_value.is_some() {
-            std::iter::once(first_value)
-                .chain(it)
-                .trust_my_length(capacity)
-                .collect_trusted()
-        } else {
-            it.collect()
-        }
-    };
-    debug_assert_eq!(ca.len(), capacity);
-    ca.rename(name);
-    ca
-}
-
-fn iterator_to_bool(
-    it: impl Iterator<Item = Option<bool>>,
-    init_null_count: usize,
-    first_value: Option<bool>,
-    name: PlSmallStr,
-    capacity: usize,
-) -> ChunkedArray<BooleanType> {
-    // safety: we know the iterators len
-    let mut ca: BooleanChunked = unsafe {
-        if init_null_count > 0 {
-            (0..init_null_count)
-                .map(|_| None)
-                .chain(std::iter::once(first_value))
-                .chain(it)
-                .trust_my_length(capacity)
-                .collect_trusted()
-        } else if first_value.is_some() {
-            std::iter::once(first_value)
-                .chain(it)
-                .trust_my_length(capacity)
-                .collect_trusted()
-        } else {
-            it.collect()
-        }
-    };
-    debug_assert_eq!(ca.len(), capacity);
-    ca.rename(name);
-    ca
-}
-
-fn iterator_to_object(
-    it: impl Iterator<Item = Option<ObjectValue>>,
-    init_null_count: usize,
-    first_value: Option<ObjectValue>,
-    name: PlSmallStr,
-    capacity: usize,
-) -> ObjectChunked<ObjectValue> {
-    // safety: we know the iterators len
-    let mut ca: ObjectChunked<ObjectValue> = unsafe {
-        if init_null_count > 0 {
-            (0..init_null_count)
-                .map(|_| None)
-                .chain(std::iter::once(first_value))
-                .chain(it)
-                .trust_my_length(capacity)
-                .collect_trusted()
-        } else if first_value.is_some() {
-            std::iter::once(first_value)
-                .chain(it)
-                .trust_my_length(capacity)
-                .collect_trusted()
-        } else {
-            it.collect()
-        }
-    };
-    debug_assert_eq!(ca.len(), capacity);
-    ca.rename(name);
-    ca
-}
-
-fn iterator_to_utf8(
-    it: impl Iterator<Item = Option<String>>,
-    init_null_count: usize,
-    first_value: Option<&str>,
-    name: PlSmallStr,
-    capacity: usize,
-) -> StringChunked {
-    let first_value = first_value.map(|v| v.to_string());
-
-    // safety: we know the iterators len
-    let mut ca: StringChunked = unsafe {
-        if init_null_count > 0 {
-            (0..init_null_count)
-                .map(|_| None)
-                .chain(std::iter::once(first_value))
-                .chain(it)
-                .trust_my_length(capacity)
-                .collect_trusted()
-        } else if first_value.is_some() {
-            std::iter::once(first_value)
-                .chain(it)
-                .trust_my_length(capacity)
-                .collect_trusted()
-        } else {
-            it.collect()
-        }
-    };
-    debug_assert_eq!(ca.len(), capacity);
-    ca.rename(name);
-    ca
-}
-
-fn iterator_to_list(
-    dt: &DataType,
-    it: impl Iterator<Item = Option<Series>>,
-    init_null_count: usize,
-    first_value: Option<&Series>,
-    name: PlSmallStr,
-    capacity: usize,
-) -> RbResult<ListChunked> {
-    let mut builder = get_list_builder(dt, capacity * 5, capacity, name);
-    for _ in 0..init_null_count {
-        builder.append_null()
-    }
-    builder
-        .append_opt_series(first_value)
-        .map_err(RbPolarsErr::from)?;
-    for opt_val in it {
-        match opt_val {
-            None => builder.append_null(),
-            Some(s) => {
-                if s.len() == 0 && s.dtype() != dt {
-                    builder
-                        .append_series(&Series::full_null(PlSmallStr::EMPTY, 0, dt))
-                        .unwrap()
-                } else {
-                    builder.append_series(&s).map_err(RbPolarsErr::from)?
-                }
-            }
-        }
-    }
-    Ok(builder.finish())
-}
+use crate::{RbResult, RbSeries};
+
+pub trait RbPolarsNumericType: PolarsNumericType {}
+
+impl RbPolarsNumericType for UInt8Type {}
+impl RbPolarsNumericType for UInt16Type {}
+impl RbPolarsNumericType for UInt32Type {}
+impl RbPolarsNumericType for UInt64Type {}
+impl RbPolarsNumericType for UInt128Type {}
+impl RbPolarsNumericType for Int8Type {}
+impl RbPolarsNumericType for Int16Type {}
+impl RbPolarsNumericType for Int32Type {}
+impl RbPolarsNumericType for Int64Type {}
+impl RbPolarsNumericType for Int128Type {}
+impl RbPolarsNumericType for Float16Type {}
+impl RbPolarsNumericType for Float32Type {}
+impl RbPolarsNumericType for Float64Type {}