polars-df 0.23.0 → 0.24.0

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

@@ -1,4 +1,5 @@
 use arrow::array::Array;
+use arrow::ffi;
 use arrow::ffi::{ArrowArrayStream, ArrowArrayStreamReader};
 use magnus::Value;
 use magnus::prelude::*;
@@ -9,6 +10,22 @@ use super::RbSeries;
 use crate::RbResult;
 use crate::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> {
     let capsule = ob.funcall("arrow_c_stream", ())?;

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

@@ -1,9 +1,10 @@
-use magnus::Value;
+use magnus::{Ruby, Value};
 
 use super::RbSeries;
 use crate::map::check_nested_object;
 use crate::map::series::{ApplyLambda, call_lambda_and_extract};
 use crate::prelude::*;
+use crate::utils::RubyAttach;
 use crate::{RbPolarsErr, RbResult};
 use crate::{apply_method_all_arrow_series2, raise_err};
 
@@ -14,7 +15,7 @@ impl RbSeries {
         return_dtype: Option<Wrap<DataType>>,
         skip_nulls: bool,
     ) -> RbResult<Self> {
-        let series = &self.series.borrow();
+        let series = &self.series.read().clone(); // Clone so we don't deadlock on re-entrance.
 
         if return_dtype.is_none() {
             polars_warn!(
@@ -58,170 +59,187 @@ impl RbSeries {
 
         }
 
-        if matches!(
-            series.dtype(),
-            DataType::Datetime(_, _)
-                | DataType::Date
-                | DataType::Duration(_)
-                | DataType::Categorical(_, _)
-                | DataType::Enum(_, _)
-                | DataType::Binary
-                | DataType::Array(_, _)
-                | DataType::Time
-                | DataType::Decimal(_, _)
-        ) || !skip_nulls
-        {
-            let mut avs = Vec::with_capacity(series.len());
-            let s = series.rechunk();
+        Ruby::attach(|_rb| {
+            if matches!(
+                series.dtype(),
+                DataType::Datetime(_, _)
+                    | DataType::Date
+                    | DataType::Duration(_)
+                    | DataType::Categorical(_, _)
+                    | DataType::Enum(_, _)
+                    | DataType::Binary
+                    | DataType::Array(_, _)
+                    | DataType::Time
+                    | DataType::Decimal(_, _)
+            ) || !skip_nulls
+            {
+                let mut avs = Vec::with_capacity(series.len());
+                let s = series.rechunk();
 
-            for av in s.iter() {
-                let out = match (skip_nulls, av) {
-                    (true, AnyValue::Null) => AnyValue::Null,
-                    (_, av) => {
-                        let av: Option<Wrap<AnyValue>> =
-                            call_lambda_and_extract(function, Wrap(av))?;
-                        match av {
-                            None => AnyValue::Null,
-                            Some(av) => av.0,
+                for av in s.iter() {
+                    let out = match (skip_nulls, av) {
+                        (true, AnyValue::Null) => AnyValue::Null,
+                        (_, av) => {
+                            let av: Option<Wrap<AnyValue>> =
+                                call_lambda_and_extract(function, Wrap(av))?;
+                            match av {
+                                None => AnyValue::Null,
+                                Some(av) => av.0,
+                            }
                         }
-                    }
-                };
-                avs.push(out)
-            }
-            let out = Series::new(series.name().clone(), &avs);
-            let dtype = out.dtype();
-            if dtype.is_nested() {
-                check_nested_object(dtype)?;
-            }
-
-            return Ok(out.into());
-        }
+                    };
+                    avs.push(out)
+                }
+                let out = Series::new(series.name().clone(), &avs);
+                let dtype = out.dtype();
+                if dtype.is_nested() {
+                    check_nested_object(dtype)?;
+                }
 
-        let out = match return_dtype {
-            Some(DataType::Int8) => {
-                let ca: Int8Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Int16) => {
-                let ca: Int16Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Int32) => {
-                let ca: Int32Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
+                return Ok(out.into());
             }
-            Some(DataType::Int64) => {
-                let ca: Int64Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::UInt8) => {
-                let ca: UInt8Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::UInt16) => {
-                let ca: UInt16Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::UInt32) => {
-                let ca: UInt32Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::UInt64) => {
-                let ca: UInt64Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Float32) => {
-                let ca: Float32Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Float64) => {
-                let ca: Float64Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    function,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Boolean) => {
-                let ca: BooleanChunked =
-                    dispatch_apply!(series, apply_lambda_with_bool_out_type, function, 0, None)?;
-                ca.into_series()
-            }
-            Some(DataType::String) => {
-                let ca =
-                    dispatch_apply!(series, apply_lambda_with_utf8_out_type, function, 0, None)?;
 
-                ca.into_series()
-            }
-            Some(DataType::List(_inner)) => {
-                todo!()
-            }
-            Some(DataType::Object(_)) => {
-                let ca =
-                    dispatch_apply!(series, apply_lambda_with_object_out_type, function, 0, None)?;
-                ca.into_series()
-            }
-            None => return dispatch_apply!(series, apply_lambda_unknown, function),
+            let out = match return_dtype {
+                Some(DataType::Int8) => {
+                    let ca: Int8Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::Int16) => {
+                    let ca: Int16Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::Int32) => {
+                    let ca: Int32Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::Int64) => {
+                    let ca: Int64Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::UInt8) => {
+                    let ca: UInt8Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::UInt16) => {
+                    let ca: UInt16Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::UInt32) => {
+                    let ca: UInt32Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::UInt64) => {
+                    let ca: UInt64Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::Float32) => {
+                    let ca: Float32Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::Float64) => {
+                    let ca: Float64Chunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_primitive_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::Boolean) => {
+                    let ca: BooleanChunked = dispatch_apply!(
+                        series,
+                        apply_lambda_with_bool_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                Some(DataType::String) => {
+                    let ca = dispatch_apply!(
+                        series,
+                        apply_lambda_with_utf8_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+
+                    ca.into_series()
+                }
+                Some(DataType::List(_inner)) => {
+                    todo!()
+                }
+                Some(DataType::Object(_)) => {
+                    let ca = dispatch_apply!(
+                        series,
+                        apply_lambda_with_object_out_type,
+                        function,
+                        0,
+                        None
+                    )?;
+                    ca.into_series()
+                }
+                None => return dispatch_apply!(series, apply_lambda_unknown, function),
 
-            _ => return dispatch_apply!(series, apply_lambda_unknown, function),
-        };
+                _ => return dispatch_apply!(series, apply_lambda_unknown, function),
+            };
 
-        Ok(RbSeries::new(out))
+            Ok(RbSeries::new(out))
+        })
     }
 }

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

@@ -8,28 +8,38 @@ mod import;
 mod map;
 mod scatter;
 
+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::{ObjectValue, RbResult};
 
 #[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),
         }
     }
 }
@@ -37,17 +47,23 @@ 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 {
-    Ruby::get().unwrap().ary_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) {
@@ -66,7 +82,7 @@ impl DataTypeFunctions for RbSeries {
         // 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 Ok(s) = &self.series.try_borrow() {
+        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

@@ -2,18 +2,21 @@ use arrow::array::Array;
 use polars::prelude::*;
 
 use crate::error::RbPolarsErr;
-use crate::{RbErr, RbResult, RbSeries};
+use crate::utils::EnterPolarsExt;
+use crate::{RbErr, 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(|| {
+            let mut s = self_.series.write();
+            match scatter(s.clone(), &idx.series.read(), &values.series.read()) {
+                Ok(out) => {
+                    *s = out;
+                    Ok(())
+                }
+                Err(e) => Err(RbErr::from(RbPolarsErr::from(e))),
             }
-            Err(e) => Err(RbErr::from(RbPolarsErr::from(e))),
-        }
+        })
     }
 }
 

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())
+}