polars-df 0.15.0 → 0.26.0

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

@@ -1,13 +1,30 @@
-use magnus::prelude::*;
+use arrow::array::Array;
+use arrow::ffi;
+use arrow::ffi::{ArrowArrayStream, ArrowArrayStreamReader};
 use magnus::Value;
-use polars::export::arrow::array::Array;
-use polars::export::arrow::ffi::{ArrowArrayStream, ArrowArrayStreamReader};
+use magnus::prelude::*;
 use polars::prelude::*;
 
 use super::RbSeries;
 
-use crate::exceptions::RbValueError;
 use crate::RbResult;
+use crate::ruby::exceptions::RbValueError;
+
+pub(crate) fn import_schema_rbcapsule(schema_capsule: Value) -> RbResult<arrow::datatypes::Field> {
+    let capsule_pointer: usize = schema_capsule.funcall("to_i", ())?;
+
+    // # Safety
+    // schema_capsule holds a valid C ArrowSchema pointer, as defined by the Arrow PyCapsule
+    // Interface
+    unsafe {
+        let schema_ptr = (capsule_pointer as *const ffi::ArrowSchema)
+            .as_ref()
+            .unwrap();
+        let field = ffi::import_field_from_c(schema_ptr).unwrap();
+
+        Ok(field)
+    }
+}
 
 /// Import `arrow_c_stream` across Ruby boundary.
 fn call_arrow_c_stream(ob: Value) -> RbResult<Value> {
@@ -39,7 +56,7 @@ pub(crate) fn import_stream_rbcapsule(capsule: Value) -> RbResult<RbSeries> {
 
     // Series::try_from fails for an empty vec of chunks
     let s = if produced_arrays.is_empty() {
-        let polars_dt = DataType::from_arrow(stream.field().dtype(), false);
+        let polars_dt = DataType::from_arrow_field(stream.field());
         Series::new_empty(stream.field().name.clone(), &polars_dt)
     } else {
         Series::try_from((stream.field(), produced_arrays)).unwrap()

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

@@ -0,0 +1,103 @@
+use magnus::{IntoValue, Ruby, Value, value::Opaque, value::ReprValue};
+
+use super::RbSeries;
+use crate::map::series::ApplyLambdaGeneric;
+use crate::prelude::*;
+use crate::ruby::gvl::GvlExt;
+use crate::ruby::utils::TryIntoValue;
+use crate::series::construction::series_from_objects;
+use crate::{RbPolarsErr, RbResult};
+use crate::{apply_all_polars_dtypes, raise_err};
+
+impl RbSeries {
+    pub fn map_elements(
+        &self,
+        function: Value,
+        return_dtype: Option<Wrap<DataType>>,
+        skip_nulls: bool,
+    ) -> RbResult<Self> {
+        let series = &self.series.read().clone(); // Clone so we don't deadlock on re-entrance.
+
+        if return_dtype.is_none() {
+            polars_warn!(
+                MapWithoutReturnDtypeWarning,
+                "Calling `map_elements` without specifying `return_dtype` can lead to unpredictable results. \
+                Specify `return_dtype` to silence this warning."
+            )
+        }
+
+        if skip_nulls && (series.null_count() == series.len()) {
+            if let Some(return_dtype) = return_dtype {
+                return Ok(
+                    Series::full_null(series.name().clone(), series.len(), &return_dtype.0).into(),
+                );
+            }
+            let msg = "The output type of the 'map_elements' function cannot be determined.\n\
+            The function was never called because 'skip_nulls: true' and all values are null.\n\
+            Consider setting 'skip_nulls: false' or setting the 'return_dtype'.";
+            raise_err!(msg, ComputeError)
+        }
+
+        let return_dtype = return_dtype.map(|dt| dt.0);
+
+        Ruby::attach(|rb| {
+            let s = match &return_dtype {
+                Some(DataType::Object(_)) => {
+                    // If the return dtype is Object we should not go through AnyValue.
+                    call_and_collect_objects(
+                        rb,
+                        series.name().clone(),
+                        function,
+                        series.len(),
+                        series.iter().map(|av| av.null_to_none().map(Wrap)),
+                        skip_nulls,
+                    )
+                }
+                Some(return_dtype) => {
+                    apply_all_polars_dtypes!(
+                        series,
+                        apply_generic_with_dtype,
+                        rb,
+                        function,
+                        return_dtype,
+                        skip_nulls
+                    )
+                }
+                None => apply_all_polars_dtypes!(series, apply_generic, rb, function, skip_nulls),
+            };
+            s.map(RbSeries::from)
+        })
+    }
+}
+
+fn call_and_collect_objects<T, I>(
+    rb: &Ruby,
+    name: PlSmallStr,
+    lambda: Value,
+    len: usize,
+    iter: I,
+    skip_nulls: bool,
+) -> RbResult<Series>
+where
+    T: TryIntoValue,
+    I: Iterator<Item = Option<T>>,
+{
+    let mut objects = Vec::with_capacity(len);
+    for opt_val in iter {
+        let arg = match opt_val {
+            None if skip_nulls => {
+                objects.push(ObjectValue {
+                    inner: Opaque::from(rb.qnil().into_value_with(rb)),
+                });
+                continue;
+            }
+            None => rb.qnil().into_value_with(rb),
+            Some(val) => val.try_into_value_with(rb)?,
+        };
+        let out: Value = lambda.funcall("call", (arg,))?;
+        objects.push(ObjectValue {
+            inner: Opaque::from(out),
+        });
+    }
+    Ok(series_from_objects(rb, name, objects))
+}

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

@@ -1,33 +1,45 @@
 mod aggregation;
 mod arithmetic;
 mod comparison;
-mod construction;
+pub(crate) mod construction;
 mod export;
 mod general;
 mod import;
+mod map;
 mod scatter;
 
-use magnus::{prelude::*, RArray};
+pub(crate) use import::import_schema_rbcapsule;
+
+use magnus::{DataTypeFunctions, RArray, Ruby, TypedData, gc, prelude::*};
+use parking_lot::RwLock;
 use polars::prelude::*;
-use std::cell::RefCell;
 
-use crate::RbResult;
+use crate::{ObjectValue, RbResult};
 
-#[magnus::wrap(class = "Polars::RbSeries")]
+#[derive(TypedData)]
+#[magnus(class = "Polars::RbSeries", mark)]
 pub struct RbSeries {
-    pub series: RefCell<Series>,
+    pub series: RwLock<Series>,
+}
+
+impl Clone for RbSeries {
+    fn clone(&self) -> Self {
+        Self {
+            series: RwLock::new(self.series.read().clone()),
+        }
+    }
 }
 
 impl From<Series> for RbSeries {
     fn from(series: Series) -> Self {
-        RbSeries::new(series)
+        Self::new(series)
     }
 }
 
 impl RbSeries {
     pub fn new(series: Series) -> Self {
         RbSeries {
-            series: RefCell::new(series),
+            series: RwLock::new(series),
         }
     }
 }
@@ -35,15 +47,45 @@ impl RbSeries {
 pub fn to_series(rs: RArray) -> RbResult<Vec<Series>> {
     let mut series = Vec::new();
     for item in rs.into_iter() {
-        series.push(<&RbSeries>::try_convert(item)?.series.borrow().clone());
+        series.push(<&RbSeries>::try_convert(item)?.series.read().clone());
     }
     Ok(series)
 }
 
-pub fn to_rbseries(s: Vec<Column>) -> RArray {
-    RArray::from_iter(
-        s.into_iter()
-            .map(|c| c.take_materialized_series())
-            .map(RbSeries::new),
-    )
+pub(crate) trait ToRbSeries {
+    fn to_rbseries(self, rb: &Ruby) -> RArray;
+}
+
+impl ToRbSeries for Vec<Column> {
+    fn to_rbseries(self, rb: &Ruby) -> RArray {
+        rb.ary_from_iter(
+            self.into_iter()
+                .map(|c| c.take_materialized_series())
+                .map(RbSeries::new),
+        )
+    }
+}
+
+pub fn mark_series(marker: &gc::Marker, series: &Series) {
+    if let DataType::Object(_) = series.dtype() {
+        for i in 0..series.len() {
+            let obj: Option<&ObjectValue> = series.get_object(i).map(|any| any.into());
+            if let Some(o) = obj {
+                marker.mark(o.inner);
+            }
+        }
+    }
+}
+
+impl DataTypeFunctions for RbSeries {
+    fn mark(&self, marker: &gc::Marker) {
+        // this is not ideal, as objects will not be marked if unable to borrow
+        // this should never happen, but log for now to avoid panic,
+        // as most series will not use Object datatype
+        if let Some(s) = &self.series.try_read() {
+            mark_series(marker, s);
+        } else {
+            eprintln!("[polars] Could not borrow!");
+        }
+    }
 }

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

@@ -1,107 +1,164 @@
-use polars::export::arrow::array::Array;
+use arrow::array::Array;
 use polars::prelude::*;
 
-use crate::error::RbPolarsErr;
-use crate::{RbErr, RbResult, RbSeries};
+use crate::utils::EnterPolarsExt;
+use crate::{RbResult, RbSeries, Ruby};
 
 impl RbSeries {
-    pub fn scatter(&self, idx: &RbSeries, values: &RbSeries) -> RbResult<()> {
-        let mut s = self.series.borrow_mut();
-        match scatter(s.clone(), &idx.series.borrow(), &values.series.borrow()) {
-            Ok(out) => {
-                *s = out;
-                Ok(())
+    pub fn scatter(rb: &Ruby, self_: &Self, idx: &RbSeries, values: &RbSeries) -> RbResult<()> {
+        rb.enter_polars(|| {
+            // We take the value because we want a ref count of 1 so that we can
+            // have mutable access cheaply via _get_inner_mut().
+            let mut lock = self_.series.write();
+            let s = std::mem::take(&mut *lock);
+            let result = scatter(
+                s,
+                &idx.clone().series.into_inner(),
+                &values.clone().series.into_inner(),
+            );
+            match result {
+                Ok(out) => {
+                    *lock = out;
+                    Ok(())
+                }
+                Err((s, e)) => {
+                    *lock = s; // Restore original series.
+                    Err(e)
+                }
             }
-            Err(e) => Err(RbErr::from(RbPolarsErr::from(e))),
-        }
+        })
     }
 }
 
-fn scatter(mut s: Series, idx: &Series, values: &Series) -> PolarsResult<Series> {
+fn scatter(s: Series, idx: &Series, values: &Series) -> Result<Series, (Series, PolarsError)> {
     let logical_dtype = s.dtype().clone();
-    let idx = idx.cast(&IDX_DTYPE)?;
-    let idx = idx.rechunk();
-    let idx = idx.idx().unwrap();
-    let idx = idx.downcast_iter().next().unwrap();
+    let converted_values;
+    let values = if logical_dtype.is_categorical() || logical_dtype.is_enum() {
+        if matches!(
+            values.dtype(),
+            DataType::Categorical(_, _) | DataType::Enum(_, _) | DataType::String | DataType::Null
+        ) {
+            converted_values = values.strict_cast(&logical_dtype);
+            match converted_values {
+                Ok(ref values) => values,
+                Err(err) => return Err((s, err)),
+            }
+        } else {
+            return Err((
+                s,
+                polars_err!(InvalidOperation: "invalid values dtype '{}' for scattering into dtype '{}'", values.dtype(), logical_dtype),
+            ));
+        }
+    } else if logical_dtype.is_decimal() {
+        if values.dtype().is_numeric() {
+            converted_values = values.strict_cast(&logical_dtype);
+            match converted_values {
+                Ok(ref values) => values,
+                Err(err) => return Err((s, err)),
+            }
+        } else {
+            return Err((
+                s,
+                polars_err!(InvalidOperation: "invalid values dtype '{}' for scattering into dtype '{}'", values.dtype(), logical_dtype),
+            ));
+        }
+    } else {
+        values
+    };
 
-    if idx.null_count() > 0 {
-        return Err(PolarsError::ComputeError(
-            "index values should not be null".into(),
+    let null_on_oob = false;
+    let idx = match polars_ops::prelude::convert_and_bound_index(idx, s.len(), null_on_oob) {
+        Ok(idx) => idx,
+        Err(err) => return Err((s, err)),
+    };
+    let idx = idx.rechunk();
+    let idx = idx.downcast_as_array();
+    if idx.has_nulls() {
+        return Err((
+            s,
+            PolarsError::ComputeError("index values should not be null".into()),
         ));
     }
-
     let idx = idx.values().as_slice();
 
-    let values = values.to_physical_repr().cast(&s.dtype().to_physical())?;
+    let mut values = match values.to_physical_repr().cast(&s.dtype().to_physical()) {
+        Ok(values) => values,
+        Err(err) => return Err((s, err)),
+    };
+
+    // Broadcast values input.
+    if values.len() == 1 && idx.len() > 1 {
+        values = values.new_from_index(0, idx.len());
+    }
+
+    let mut phys = s.to_physical_repr().into_owned();
+    drop(s); // Reduce refcount to make use of in-place mutation of possible.
+    let ret = scatter_impl(&mut phys, &logical_dtype, idx, &values);
+    match ret {
+        Ok(s) => Ok(unsafe { s.from_physical_unchecked(&logical_dtype).unwrap() }),
+        Err(e) => Err((
+            unsafe { phys.from_physical_unchecked(&logical_dtype).unwrap() },
+            e,
+        )),
+    }
+}
 
-    // do not shadow, otherwise s is not dropped immediately
-    // and we want to have mutable access
-    s = s.to_physical_repr().into_owned();
+macro_rules! with_match_physical_numeric_polars_type {(
+    $key_type:expr, | $_:tt $T:ident | $($body:tt)*
+) => ({
+    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
+    use polars_core::datatypes::DataType::*;
+    match $key_type {
+        Int8 => __with_ty__! { Int8Type },
+        Int16 => __with_ty__! { Int16Type },
+        Int32 => __with_ty__! { Int32Type },
+        Int64 => __with_ty__! { Int64Type },
+        Int128 => __with_ty__! { Int128Type },
+        UInt8 => __with_ty__! { UInt8Type },
+        UInt16 => __with_ty__! { UInt16Type },
+        UInt32 => __with_ty__! { UInt32Type },
+        UInt64 => __with_ty__! { UInt64Type },
+        UInt128 => __with_ty__! { UInt128Type },
+        Float16 => __with_ty__! { Float16Type },
+        Float32 => __with_ty__! { Float32Type },
+        Float64 => __with_ty__! { Float64Type },
+        dt => panic!("not implemented for dtype {:?}", dt),
+    }
+})}
+
+fn scatter_impl(
+    s: &mut Series,
+    logical_dtype: &DataType,
+    idx: &[IdxSize],
+    values: &Series,
+) -> PolarsResult<Series> {
     let mutable_s = s._get_inner_mut();
 
-    let s = match logical_dtype.to_physical() {
-        DataType::Int8 => {
-            let ca: &mut ChunkedArray<Int8Type> = mutable_s.as_mut();
-            let values = values.i8()?;
-            std::mem::take(ca).scatter(idx, values)
-        }
-        DataType::Int16 => {
-            let ca: &mut ChunkedArray<Int16Type> = mutable_s.as_mut();
-            let values = values.i16()?;
-            std::mem::take(ca).scatter(idx, values)
-        }
-        DataType::Int32 => {
-            let ca: &mut ChunkedArray<Int32Type> = mutable_s.as_mut();
-            let values = values.i32()?;
-            std::mem::take(ca).scatter(idx, values)
-        }
-        DataType::Int64 => {
-            let ca: &mut ChunkedArray<Int64Type> = mutable_s.as_mut();
-            let values = values.i64()?;
-            std::mem::take(ca).scatter(idx, values)
-        }
-        DataType::UInt8 => {
-            let ca: &mut ChunkedArray<UInt8Type> = mutable_s.as_mut();
-            let values = values.u8()?;
-            std::mem::take(ca).scatter(idx, values)
-        }
-        DataType::UInt16 => {
-            let ca: &mut ChunkedArray<UInt16Type> = mutable_s.as_mut();
-            let values = values.u16()?;
-            std::mem::take(ca).scatter(idx, values)
-        }
-        DataType::UInt32 => {
-            let ca: &mut ChunkedArray<UInt32Type> = mutable_s.as_mut();
-            let values = values.u32()?;
-            std::mem::take(ca).scatter(idx, values)
-        }
-        DataType::UInt64 => {
-            let ca: &mut ChunkedArray<UInt64Type> = mutable_s.as_mut();
-            let values = values.u64()?;
-            std::mem::take(ca).scatter(idx, values)
-        }
-        DataType::Float32 => {
-            let ca: &mut ChunkedArray<Float32Type> = mutable_s.as_mut();
-            let values = values.f32()?;
-            std::mem::take(ca).scatter(idx, values)
-        }
-        DataType::Float64 => {
-            let ca: &mut ChunkedArray<Float64Type> = mutable_s.as_mut();
-            let values = values.f64()?;
-            std::mem::take(ca).scatter(idx, values)
+    match mutable_s.dtype() {
+        dt if dt.is_primitive_numeric() => {
+            with_match_physical_numeric_polars_type!(dt, |$T| {
+                let ca: &mut ChunkedArray<$T> = mutable_s.as_mut();
+                let values: &ChunkedArray<$T> = values.as_ref().as_ref();
+                ca.scatter(idx, values.iter())
+            })
         }
         DataType::Boolean => {
-            let ca = s.bool()?;
+            let ca: &mut ChunkedArray<BooleanType> = mutable_s.as_mut();
             let values = values.bool()?;
-            ca.scatter(idx, values)
+            ca.scatter(idx, values.iter())
+        }
+        DataType::Binary => {
+            let ca: &mut ChunkedArray<BinaryType> = mutable_s.as_mut();
+            let values = values.binary()?;
+            ca.scatter(idx, values.iter())
         }
         DataType::String => {
-            let ca = s.str()?;
+            let ca: &mut ChunkedArray<StringType> = mutable_s.as_mut();
             let values = values.str()?;
-            ca.scatter(idx, values)
+            ca.scatter(idx, values.iter())
         }
-        _ => panic!("not yet implemented for dtype: {}", logical_dtype),
-    };
-
-    s.and_then(|s| s.cast(&logical_dtype))
+        _ => Err(PolarsError::ComputeError(
+            format!("not yet implemented for dtype: {logical_dtype}").into(),
+        )),
+    }
 }

data/ext/polars/src/sql.rs

@@ -1,13 +1,20 @@
+use parking_lot::RwLock;
 use polars::sql::SQLContext;
-use std::cell::RefCell;
 
 use crate::{RbLazyFrame, RbPolarsErr, RbResult};
 
 #[magnus::wrap(class = "Polars::RbSQLContext")]
 #[repr(transparent)]
-#[derive(Clone)]
 pub struct RbSQLContext {
-    pub context: RefCell<SQLContext>,
+    pub context: RwLock<SQLContext>,
+}
+
+impl Clone for RbSQLContext {
+    fn clone(&self) -> Self {
+        Self {
+            context: RwLock::new(self.context.read().clone()),
+        }
+    }
 }
 
 #[allow(
@@ -19,30 +26,30 @@ impl RbSQLContext {
     #[allow(clippy::new_without_default)]
     pub fn new() -> RbSQLContext {
         RbSQLContext {
-            context: SQLContext::new().into(),
+            context: RwLock::new(SQLContext::new()),
         }
     }
 
     pub fn execute(&self, query: String) -> RbResult<RbLazyFrame> {
         Ok(self
             .context
-            .borrow_mut()
+            .write()
             .execute(&query)
             .map_err(RbPolarsErr::from)?
             .into())
     }
 
     pub fn get_tables(&self) -> RbResult<Vec<String>> {
-        Ok(self.context.borrow().get_tables())
+        Ok(self.context.read().get_tables())
     }
 
     pub fn register(&self, name: String, lf: &RbLazyFrame) {
         self.context
-            .borrow_mut()
-            .register(&name, lf.ldf.borrow().clone())
+            .write()
+            .register(&name, lf.clone().ldf.into_inner())
     }
 
     pub fn unregister(&self, name: String) {
-        self.context.borrow_mut().unregister(&name)
+        self.context.write().unregister(&name)
     }
 }

data/ext/polars/src/testing/frame.rs

@@ -0,0 +1,31 @@
+use polars_testing::asserts::{DataFrameEqualOptions, assert_dataframe_equal};
+
+use crate::error::RbPolarsErr;
+use crate::{RbDataFrame, RbResult};
+
+pub fn assert_dataframe_equal_rb(
+    left: &RbDataFrame,
+    right: &RbDataFrame,
+    check_row_order: bool,
+    check_column_order: bool,
+    check_dtypes: bool,
+    check_exact: bool,
+    rel_tol: f64,
+    abs_tol: f64,
+    categorical_as_str: bool,
+) -> RbResult<()> {
+    let left_df = &left.df.read();
+    let right_df = &right.df.read();
+
+    let options = DataFrameEqualOptions {
+        check_row_order,
+        check_column_order,
+        check_dtypes,
+        check_exact,
+        rel_tol,
+        abs_tol,
+        categorical_as_str,
+    };
+
+    assert_dataframe_equal(left_df, right_df, options).map_err(|e| RbPolarsErr::from(e).into())
+}

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

@@ -0,0 +1,5 @@
+mod frame;
+mod series;
+
+pub use frame::*;
+pub use series::*;

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

@@ -0,0 +1,31 @@
+use polars_testing::asserts::{SeriesEqualOptions, assert_series_equal};
+
+use crate::error::RbPolarsErr;
+use crate::{RbResult, RbSeries};
+
+pub fn assert_series_equal_rb(
+    left: &RbSeries,
+    right: &RbSeries,
+    check_dtypes: bool,
+    check_names: bool,
+    check_order: bool,
+    check_exact: bool,
+    rel_tol: f64,
+    abs_tol: f64,
+    categorical_as_str: bool,
+) -> RbResult<()> {
+    let left_series = &left.series.read();
+    let right_series = &right.series.read();
+
+    let options = SeriesEqualOptions {
+        check_dtypes,
+        check_names,
+        check_order,
+        check_exact,
+        rel_tol,
+        abs_tol,
+        categorical_as_str,
+    };
+
+    assert_series_equal(left_series, right_series, options).map_err(|e| RbPolarsErr::from(e).into())
+}