polars-df 0.15.0 → 0.26.0

data/ext/polars/src/on_startup.rs

@@ -1,32 +1,215 @@
+#![allow(unsafe_op_in_unsafe_fn)]
 use std::any::Any;
 use std::sync::Arc;
+use std::sync::OnceLock;
 
-use magnus::IntoValue;
+use arrow::array::Array;
+use magnus::{IntoValue, Ruby, Value, prelude::*, value::Opaque};
+use polars::chunked_array::object::ObjectArray;
 use polars::prelude::*;
 use polars_core::chunked_array::object::builder::ObjectChunkedBuilder;
 use polars_core::chunked_array::object::registry;
 use polars_core::chunked_array::object::registry::AnonymousObjectBuilder;
 use polars_core::prelude::AnyValue;
+use polars_error::PolarsWarning;
+use polars_error::signals::register_polars_keyboard_interrupt_hook;
 
-use crate::prelude::ObjectValue;
 use crate::Wrap;
+use crate::dataframe::RbDataFrame;
+use crate::lazyframe::RbLazyFrame;
+use crate::map::lazy::call_lambda_with_series;
+use crate::prelude::ObjectValue;
+use crate::rb_modules::pl_utils;
+use crate::ruby::gvl::GvlExt;
+use crate::ruby::ruby_convert_registry::{FromRubyConvertRegistry, RubyConvertRegistry};
+use crate::ruby::ruby_udf;
+use crate::ruby::thread::{is_non_ruby_thread, run_in_ruby_thread};
+use crate::ruby::utils::{TryIntoValue, to_pl_err};
+use crate::series::RbSeries;
+
+fn ruby_function_caller_series(
+    s: &[Column],
+    output_dtype: Option<DataType>,
+    lambda: Opaque<Value>,
+) -> PolarsResult<Column> {
+    if is_non_ruby_thread() {
+        let s2 = s.to_vec();
+        return run_in_ruby_thread(move || ruby_function_caller_series(&s2, output_dtype, lambda));
+    }
+
+    Ruby::attach(|rb| call_lambda_with_series(rb, s, output_dtype, lambda))
+}
+
+fn ruby_function_caller_df(df: DataFrame, lambda: Opaque<Value>) -> PolarsResult<DataFrame> {
+    Ruby::attach(|rb| {
+        let lambda = rb.get_inner(lambda);
+
+        let rbpolars = pl_utils(rb);
+
+        // create a RbSeries struct/object for Ruby
+        let rbdf = RbDataFrame::new(df);
+        // Wrap this RbSeries object in the Ruby side Series wrapper
+        let ruby_df_wrapper: Value = rbpolars.funcall("wrap_df", (rbdf,)).map_err(to_pl_err)?;
+
+        // call the lambda and get a Ruby side df wrapper
+        let result_df_wrapper: Value = lambda
+            .funcall("call", (ruby_df_wrapper,))
+            .map_err(to_pl_err)?;
+
+        // unpack the wrapper in a RbDataFrame
+        let rbdf: &RbDataFrame = result_df_wrapper.funcall("_df", ()).map_err(|_| {
+            let rbtype = unsafe { result_df_wrapper.classname() };
+            PolarsError::ComputeError(
+                format!("Expected 'LazyFrame.map' to return a 'DataFrame', got a '{rbtype}'",)
+                    .into(),
+            )
+        })?;
+        Ok(rbdf.clone().df.into_inner())
+    })
+}
+
+fn warning_function(msg: &str, _warning: PolarsWarning) {
+    Ruby::attach(|rb| {
+        if let Err(e) = pl_utils(rb).funcall::<_, _, Value>("_polars_warn", (msg.to_string(),)) {
+            eprintln!("{e}")
+        }
+    })
+}
 
-pub(crate) fn register_startup_deps() {
-    if !registry::is_object_builder_registered() {
+static POLARS_REGISTRY_INIT_LOCK: OnceLock<()> = OnceLock::new();
+
+pub unsafe fn register_startup_deps(catch_keyboard_interrupt: bool) {
+    POLARS_REGISTRY_INIT_LOCK.get_or_init(|| {
         let object_builder = Box::new(|name: PlSmallStr, capacity: usize| {
             Box::new(ObjectChunkedBuilder::<ObjectValue>::new(name, capacity))
                 as Box<dyn AnonymousObjectBuilder>
         });
 
         let object_converter = Arc::new(|av: AnyValue| {
-            let object = ObjectValue {
-                inner: Wrap(av).into_value().into(),
-            };
+            let object = Ruby::attach(|rb| ObjectValue {
+                inner: Wrap(av).try_into_value_with(rb).unwrap().into(),
+            });
+            Box::new(object) as Box<dyn Any>
+        });
+        let rbobject_converter = Arc::new(|av: AnyValue| {
+            let object = Ruby::attach(|rb| Wrap(av).try_into_value_with(rb).unwrap());
             Box::new(object) as Box<dyn Any>
         });
+        fn object_array_getter(arr: &dyn Array, idx: usize) -> Option<AnyValue<'_>> {
+            let arr = arr
+                .as_any()
+                .downcast_ref::<ObjectArray<ObjectValue>>()
+                .unwrap();
+            arr.get(idx).map(|v| AnyValue::Object(v))
+        }
+        fn with_gil(f: &mut dyn FnMut()) {
+            Ruby::attach(|_| f())
+        }
+
+        crate::ruby::ruby_convert_registry::register_converters(RubyConvertRegistry {
+            from_rb: FromRubyConvertRegistry {
+                file_provider_result: Arc::new(|_rb_f| Ruby::attach(|_rb| todo!())),
+                series: Arc::new(|rb_f| {
+                    Ruby::attach(|rb| {
+                        Ok(Box::new(
+                            <&RbSeries>::try_convert(rb.get_inner(rb_f))?
+                                .clone()
+                                .series
+                                .into_inner(),
+                        ) as _)
+                    })
+                }),
+                df: Arc::new(|rb_f| {
+                    Ruby::attach(|rb| {
+                        Ok(Box::new(
+                            <&RbDataFrame>::try_convert(rb.get_inner(rb_f))?
+                                .clone()
+                                .df
+                                .into_inner(),
+                        ) as _)
+                    })
+                }),
+                dsl_plan: Arc::new(|rb_f| {
+                    Ruby::attach(|rb| {
+                        Ok(Box::new(
+                            <&RbLazyFrame>::try_convert(rb.get_inner(rb_f))?
+                                .clone()
+                                .ldf
+                                .into_inner()
+                                .logical_plan,
+                        ) as _)
+                    })
+                }),
+                schema: Arc::new(|rb_f| {
+                    Ruby::attach(|rb| {
+                        Ok(Box::new(
+                            Wrap::<polars_core::schema::Schema>::try_convert(rb.get_inner(rb_f))?.0,
+                        ) as _)
+                    })
+                }),
+            },
+            to_rb: crate::ruby::ruby_convert_registry::ToRubyConvertRegistry {
+                df: Arc::new(|df| {
+                    Ruby::attach(|rb| {
+                        Ok(
+                            RbDataFrame::new(df.downcast_ref::<DataFrame>().unwrap().clone())
+                                .into_value_with(rb),
+                        )
+                    })
+                }),
+                series: Arc::new(|series| {
+                    Ruby::attach(|rb| {
+                        Ok(
+                            RbSeries::new(series.downcast_ref::<Series>().unwrap().clone())
+                                .into_value_with(rb),
+                        )
+                    })
+                }),
+                dsl_plan: Arc::new(|dsl_plan| {
+                    Ruby::attach(|rb| {
+                        Ok(RbLazyFrame::from(LazyFrame::from(
+                            dsl_plan
+                                .downcast_ref::<polars_plan::dsl::DslPlan>()
+                                .unwrap()
+                                .clone(),
+                        ))
+                        .into_value_with(rb))
+                    })
+                }),
+                schema: Arc::new(|schema| {
+                    Ruby::attach(|rb| {
+                        Wrap(
+                            schema
+                                .downcast_ref::<polars_core::schema::Schema>()
+                                .unwrap()
+                                .clone(),
+                        )
+                        .try_into_value_with(rb)
+                    })
+                }),
+            },
+        });
 
         let object_size = std::mem::size_of::<ObjectValue>();
         let physical_dtype = ArrowDataType::FixedSizeBinary(object_size);
-        registry::register_object_builder(object_builder, object_converter, physical_dtype)
-    }
+        registry::register_object_builder(
+            object_builder,
+            object_converter,
+            rbobject_converter,
+            physical_dtype,
+            Arc::new(object_array_getter),
+            Arc::new(with_gil),
+        );
+
+        // Register SERIES UDF.
+        ruby_udf::CALL_COLUMNS_UDF_RUBY = Some(ruby_function_caller_series);
+        // Register DATAFRAME UDF.
+        ruby_udf::CALL_DF_UDF_RUBY = Some(ruby_function_caller_df);
+        // Register warning function for `polars_warn!`.
+        polars_error::set_warning_function(warning_function);
+
+        if catch_keyboard_interrupt {
+            register_polars_keyboard_interrupt_hook();
+        }
+    });
 }

data/ext/polars/src/prelude.rs

@@ -1,3 +1,4 @@
 pub use polars::prelude::*;
 
 pub use crate::conversion::*;
+pub(crate) use crate::rb_modules;

data/ext/polars/src/rb_modules.rs

@@ -1,64 +1,17 @@
-use magnus::{value::Lazy, ExceptionClass, Module, RClass, RModule, Ruby};
+use magnus::{Module, RClass, RModule, Ruby, value::Lazy};
 
-static POLARS: Lazy<RModule> = Lazy::new(|ruby| ruby.class_object().const_get("Polars").unwrap());
+static POLARS: Lazy<RModule> = Lazy::new(|rb| rb.class_object().const_get("Polars").unwrap());
+static UTILS: Lazy<RModule> = Lazy::new(|rb| rb.get_inner(&POLARS).const_get("Utils").unwrap());
+static SERIES: Lazy<RClass> = Lazy::new(|rb| rb.get_inner(&POLARS).const_get("Series").unwrap());
 
-pub(crate) fn polars() -> RModule {
-    Ruby::get().unwrap().get_inner(&POLARS)
+pub(crate) fn polars(rb: &Ruby) -> RModule {
+    rb.get_inner(&POLARS)
 }
 
-static SERIES: Lazy<RClass> =
-    Lazy::new(|ruby| ruby.get_inner(&POLARS).const_get("Series").unwrap());
-
-pub(crate) fn series() -> RClass {
-    Ruby::get().unwrap().get_inner(&SERIES)
-}
-
-static UTILS: Lazy<RModule> = Lazy::new(|ruby| ruby.get_inner(&POLARS).const_get("Utils").unwrap());
-
-pub(crate) fn utils() -> RModule {
-    Ruby::get().unwrap().get_inner(&UTILS)
-}
-
-static BIGDECIMAL: Lazy<RClass> =
-    Lazy::new(|ruby| ruby.class_object().const_get("BigDecimal").unwrap());
-
-pub(crate) fn bigdecimal() -> RClass {
-    Ruby::get().unwrap().get_inner(&BIGDECIMAL)
-}
-
-static DATE: Lazy<RClass> = Lazy::new(|ruby| ruby.class_object().const_get("Date").unwrap());
-
-pub(crate) fn date() -> RClass {
-    Ruby::get().unwrap().get_inner(&DATE)
+pub(crate) fn pl_utils(rb: &Ruby) -> RModule {
+    rb.get_inner(&UTILS)
 }
 
-static DATETIME: Lazy<RClass> =
-    Lazy::new(|ruby| ruby.class_object().const_get("DateTime").unwrap());
-
-pub(crate) fn datetime() -> RClass {
-    Ruby::get().unwrap().get_inner(&DATETIME)
-}
-
-static ERROR: Lazy<ExceptionClass> =
-    Lazy::new(|ruby| ruby.get_inner(&POLARS).const_get("Error").unwrap());
-
-pub(crate) fn error() -> ExceptionClass {
-    Ruby::get().unwrap().get_inner(&ERROR)
-}
-
-static COMPUTE_ERROR: Lazy<ExceptionClass> =
-    Lazy::new(|ruby| ruby.get_inner(&POLARS).const_get("ComputeError").unwrap());
-
-pub(crate) fn compute_error() -> ExceptionClass {
-    Ruby::get().unwrap().get_inner(&COMPUTE_ERROR)
-}
-
-static INVALID_OPERATION_ERROR: Lazy<ExceptionClass> = Lazy::new(|ruby| {
-    ruby.get_inner(&POLARS)
-        .const_get("InvalidOperationError")
-        .unwrap()
-});
-
-pub(crate) fn invalid_operation_error() -> ExceptionClass {
-    Ruby::get().unwrap().get_inner(&INVALID_OPERATION_ERROR)
+pub(crate) fn pl_series(rb: &Ruby) -> RClass {
+    rb.get_inner(&SERIES)
 }

data/ext/polars/src/ruby/exceptions.rs

@@ -0,0 +1,26 @@
+use magnus::{Error, Ruby};
+use std::borrow::Cow;
+
+macro_rules! create_ruby_exception {
+    ($type:ident, $cls:ident) => {
+        pub struct $type {}
+
+        impl $type {
+            pub fn new_err<T>(message: T) -> Error
+            where
+                T: Into<Cow<'static, str>>,
+            {
+                let cls = Ruby::get().unwrap().$cls();
+                Error::new(cls, message)
+            }
+        }
+    };
+}
+
+create_ruby_exception!(RbKeyboardInterrupt, exception_interrupt);
+create_ruby_exception!(RbIndexError, exception_index_error);
+create_ruby_exception!(RbIOError, exception_io_error);
+create_ruby_exception!(RbOverflowError, exception_range_error);
+create_ruby_exception!(RbRuntimeError, exception_runtime_error);
+create_ruby_exception!(RbTypeError, exception_type_error);
+create_ruby_exception!(RbValueError, exception_arg_error);

data/ext/polars/src/ruby/gvl.rs

@@ -0,0 +1,104 @@
+use std::ffi::c_void;
+use std::ptr::null_mut;
+
+use magnus::Ruby;
+use magnus::error::RubyUnavailableError;
+use rb_sys::{rb_thread_call_with_gvl, rb_thread_call_without_gvl};
+
+pub trait GvlExt {
+    fn attach<T, F>(func: F) -> T
+    where
+        F: FnOnce(&Ruby) -> T;
+
+    fn detach<T, F>(&self, func: F) -> T
+    where
+        F: Send + FnOnce() -> T,
+        T: Send;
+}
+
+unsafe extern "C" {
+    fn ruby_thread_has_gvl_p() -> std::ffi::c_int;
+}
+
+impl GvlExt for Ruby {
+    fn attach<T, F>(func: F) -> T
+    where
+        F: FnOnce(&Ruby) -> T,
+    {
+        // recheck GVL state since cached value can be incorrect
+        // https://github.com/matsadler/magnus/pull/161
+        if let Ok(rb) = Ruby::get()
+            && unsafe { ruby_thread_has_gvl_p() } != 0
+        {
+            func(&rb)
+        } else if !matches!(Ruby::get(), Err(RubyUnavailableError::NonRubyThread)) {
+            let mut data = CallbackData {
+                func: Some(func),
+                result: None,
+            };
+
+            unsafe {
+                rb_thread_call_with_gvl(
+                    Some(call_with_gvl::<F, T>),
+                    &mut data as *mut _ as *mut c_void,
+                );
+            }
+
+            data.result.unwrap()
+        } else {
+            panic!("Non-Ruby thread");
+        }
+    }
+
+    fn detach<T, F>(&self, func: F) -> T
+    where
+        F: Send + FnOnce() -> T,
+        T: Send,
+    {
+        if std::env::var("POLARS_GVL").is_ok() {
+            func()
+        } else {
+            let mut data = CallbackData {
+                func: Some(func),
+                result: None,
+            };
+
+            unsafe {
+                rb_thread_call_without_gvl(
+                    Some(call_without_gvl::<F, T>),
+                    &mut data as *mut _ as *mut c_void,
+                    None,
+                    null_mut(),
+                );
+            }
+
+            data.result.unwrap()
+        }
+    }
+}
+
+struct CallbackData<F, T> {
+    func: Option<F>,
+    result: Option<T>,
+}
+
+extern "C" fn call_without_gvl<F, T>(data: *mut c_void) -> *mut c_void
+where
+    F: FnOnce() -> T,
+{
+    let data = unsafe { &mut *(data as *mut CallbackData<F, T>) };
+    let func = data.func.take().unwrap();
+    data.result = Some(func());
+    null_mut()
+}
+
+extern "C" fn call_with_gvl<F, T>(data: *mut c_void) -> *mut c_void
+where
+    F: FnOnce(&Ruby) -> T,
+{
+    let rb = Ruby::get().unwrap();
+    let data = unsafe { &mut *(data as *mut CallbackData<F, T>) };
+    let func = data.func.take().unwrap();
+    data.result = Some(func(&rb));
+    null_mut()
+}

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

@@ -0,0 +1,46 @@
+use magnus::{Ruby, Value, value::Opaque};
+use polars::prelude::{AllowedOptimizations, DataFrame, LazyFrame, SchemaRef};
+
+use crate::ruby::ruby_function::RubyFunction;
+use crate::ruby::ruby_udf::CALL_DF_UDF_RUBY;
+use crate::ruby::thread::{is_non_ruby_thread, run_in_ruby_thread, start_background_ruby_thread};
+use crate::ruby::utils::ArcValue;
+
+pub trait RubyUdfLazyFrameExt {
+    fn map_ruby(
+        self,
+        function: RubyFunction,
+        optimizations: AllowedOptimizations,
+        schema: Option<SchemaRef>,
+        validate_output: bool,
+    ) -> LazyFrame;
+}
+
+impl RubyUdfLazyFrameExt for LazyFrame {
+    fn map_ruby(
+        self,
+        function: RubyFunction,
+        optimizations: AllowedOptimizations,
+        _schema: Option<SchemaRef>,
+        _validate_output: bool,
+    ) -> LazyFrame {
+        let f = move |df: DataFrame, lambda: Opaque<Value>| {
+            let func = unsafe { CALL_DF_UDF_RUBY.unwrap() };
+            func(df, lambda)
+        };
+
+        // handle non-Ruby threads
+        start_background_ruby_thread(&Ruby::get().unwrap());
+        let udf = ArcValue::new(function.0);
+        let f = move |df| {
+            if is_non_ruby_thread() {
+                let udf = udf.clone();
+                return run_in_ruby_thread(move || f(df, *udf.0));
+            }
+            f(df, *udf.0)
+        };
+
+        let schema = None;
+        self.map(f, optimizations, schema, Some("RUBY UDF"))
+    }
+}

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

@@ -0,0 +1,11 @@
+pub mod exceptions;
+pub mod gvl;
+pub mod lazy;
+pub mod numo;
+pub mod plan_callback;
+pub mod rb_modules;
+pub mod ruby_convert_registry;
+pub mod ruby_function;
+pub mod ruby_udf;
+pub mod thread;
+pub mod utils;

data/ext/polars/src/ruby/numo.rs

@@ -0,0 +1,52 @@
+use magnus::{IntoValue, Module, RClass, RModule, Ruby, Value, prelude::*};
+
+use crate::RbResult;
+
+pub trait Element: IntoValue {
+    fn class_name() -> &'static str;
+}
+
+macro_rules! create_element {
+    ($type:ty, $name:expr) => {
+        impl Element for $type {
+            fn class_name() -> &'static str {
+                $name
+            }
+        }
+    };
+}
+
+create_element!(i8, "Int8");
+create_element!(i16, "Int16");
+create_element!(i32, "Int32");
+create_element!(i64, "Int64");
+create_element!(u8, "UInt8");
+create_element!(u16, "UInt16");
+create_element!(u32, "UInt32");
+create_element!(u64, "UInt64");
+create_element!(f32, "SFloat");
+create_element!(f64, "DFloat");
+create_element!(bool, "Bit");
+
+impl<T> Element for Option<T>
+where
+    Option<T>: IntoValue,
+{
+    fn class_name() -> &'static str {
+        "RObject"
+    }
+}
+
+pub struct RbArray1<T>(T);
+
+impl<T: Element> RbArray1<T> {
+    pub fn from_iter<I>(rb: &Ruby, values: I) -> RbResult<Value>
+    where
+        I: IntoIterator<Item = T>,
+    {
+        rb.class_object()
+            .const_get::<_, RModule>("Numo")?
+            .const_get::<_, RClass>(T::class_name())?
+            .funcall("cast", (rb.ary_from_iter(values),))
+    }
+}