polars-df 0.15.0 → 0.26.0

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

@@ -0,0 +1,198 @@
+use magnus::{Ruby, Value, value::Opaque, value::ReprValue};
+use polars_plan::prelude::PlanCallback;
+
+use crate::RbResult;
+use crate::ruby::gvl::GvlExt;
+use crate::ruby::ruby_function::RubyFunction;
+use crate::ruby::thread::{is_non_ruby_thread, run_in_ruby_thread, start_background_ruby_thread};
+use crate::ruby::utils::{ArcValue, to_pl_err};
+
+pub trait PlanCallbackArgs {
+    fn into_rbany(self, rb: &Ruby) -> RbResult<Value>;
+}
+
+pub trait PlanCallbackOut: Sized {
+    fn from_rbany(rbany: Value, rb: &Ruby) -> RbResult<Self>;
+}
+
+mod _ruby {
+    use std::sync::Arc;
+
+    use crate::RbResult;
+    use magnus::{IntoValue, RArray, Ruby, TryConvert, Value, value::ReprValue};
+    use polars_utils::pl_str::PlSmallStr;
+
+    macro_rules! impl_rbcb_type {
+        ($($type:ty),+) => {
+            $(
+            impl super::PlanCallbackArgs for $type {
+                fn into_rbany(self, rb: &Ruby) -> RbResult<Value> {
+                    Ok(self.into_value_with(rb))
+                }
+            }
+
+            impl super::PlanCallbackOut for $type {
+                fn from_rbany(rbany: Value, _rb: &Ruby) -> RbResult<Self> {
+                    Self::try_convert(rbany)
+                }
+            }
+            )+
+        };
+    }
+
+    macro_rules! impl_rbcb_type_to_from {
+        ($($type:ty => $transformed:ty),+) => {
+            $(
+            impl super::PlanCallbackArgs for $type {
+                fn into_rbany(self, rb: &Ruby) -> RbResult<Value> {
+                    Ok(<$transformed>::from(self).into_value_with(rb))
+                }
+            }
+
+            impl super::PlanCallbackOut for $type {
+                fn from_rbany(rbany: Value, _rb: &Ruby) -> RbResult<Self> {
+                    <$transformed>::try_convert(rbany).map(Into::into)
+                }
+            }
+            )+
+        };
+    }
+
+    macro_rules! impl_registrycb_type {
+        ($(($type:path, $from:ident, $to:ident)),+) => {
+            $(
+            impl super::PlanCallbackArgs for $type {
+                fn into_rbany(self, _rb: &Ruby) -> RbResult<Value> {
+                    let registry = crate::ruby::ruby_convert_registry::get_ruby_convert_registry();
+                    (registry.to_rb.$to)(&self)
+                }
+            }
+
+            impl super::PlanCallbackOut for $type {
+                fn from_rbany(rbany: Value, _rb: &Ruby) -> RbResult<Self> {
+                    let registry = crate::ruby::ruby_convert_registry::get_ruby_convert_registry();
+                    // TODO remove into
+                    let obj = (registry.from_rb.$from)(rbany.into())?;
+                    let obj = obj.downcast().unwrap();
+                    Ok(*obj)
+                }
+            }
+            )+
+        };
+    }
+
+    impl<T: super::PlanCallbackArgs> super::PlanCallbackArgs for Option<T> {
+        fn into_rbany(self, rb: &Ruby) -> RbResult<Value> {
+            match self {
+                None => Ok(rb.qnil().as_value()),
+                Some(v) => v.into_rbany(rb),
+            }
+        }
+    }
+
+    impl<T: super::PlanCallbackOut> super::PlanCallbackOut for Option<T> {
+        fn from_rbany(rbany: Value, rb: &Ruby) -> RbResult<Self> {
+            if rbany.is_nil() {
+                Ok(None)
+            } else {
+                T::from_rbany(rbany, rb).map(Some)
+            }
+        }
+    }
+
+    impl<T, U> super::PlanCallbackArgs for (T, U)
+    where
+        T: super::PlanCallbackArgs,
+        U: super::PlanCallbackArgs,
+    {
+        fn into_rbany(self, rb: &Ruby) -> RbResult<Value> {
+            Ok(rb
+                .ary_new_from_values(&[self.0.into_rbany(rb)?, self.1.into_rbany(rb)?])
+                .as_value())
+        }
+    }
+
+    impl<T, U> super::PlanCallbackOut for (T, U)
+    where
+        T: super::PlanCallbackOut,
+        U: super::PlanCallbackOut,
+    {
+        fn from_rbany(rbany: Value, rb: &Ruby) -> RbResult<Self> {
+            let tuple = RArray::try_convert(rbany)?;
+            Ok((
+                T::from_rbany(tuple.entry(0)?, rb)?,
+                U::from_rbany(tuple.entry(1)?, rb)?,
+            ))
+        }
+    }
+
+    impl_rbcb_type! {
+        bool,
+        usize,
+        String
+    }
+    impl_rbcb_type_to_from! {
+        PlSmallStr => String
+    }
+    impl_registrycb_type! {
+        (polars_core::series::Series, series, series),
+        (polars_core::frame::DataFrame, df, df),
+        (polars_plan::dsl::DslPlan, dsl_plan, dsl_plan),
+        (polars_core::schema::Schema, schema, schema)
+    }
+
+    impl<T: super::PlanCallbackArgs + Clone> super::PlanCallbackArgs for Arc<T> {
+        fn into_rbany(self, rb: &Ruby) -> RbResult<Value> {
+            Arc::unwrap_or_clone(self).into_rbany(rb)
+        }
+    }
+
+    impl<T: super::PlanCallbackArgs + Clone> super::PlanCallbackArgs for Vec<T> {
+        fn into_rbany(self, rb: &Ruby) -> RbResult<Value> {
+            rb.ary_try_from_iter(self.into_iter().map(|v| v.into_rbany(rb)))
+                .map(|v| v.as_value())
+        }
+    }
+
+    impl<T: super::PlanCallbackOut> super::PlanCallbackOut for Arc<T> {
+        fn from_rbany(rbany: Value, rb: &Ruby) -> RbResult<Self> {
+            T::from_rbany(rbany, rb).map(Arc::from)
+        }
+    }
+}
+
+pub(crate) trait PlanCallbackExt<Args, Out> {
+    fn new_ruby(rbfn: RubyFunction) -> Self;
+}
+
+impl<Args: PlanCallbackArgs + Send + 'static, Out: PlanCallbackOut + Send + 'static>
+    PlanCallbackExt<Args, Out> for PlanCallback<Args, Out>
+{
+    fn new_ruby(rbfn: RubyFunction) -> Self {
+        let f = move |args: Args, lambda: Opaque<Value>| {
+            Ruby::attach(|rb| {
+                let out = Out::from_rbany(
+                    rb.get_inner(lambda)
+                        .funcall("call", (args.into_rbany(rb).map_err(to_pl_err)?,))
+                        .map_err(to_pl_err)?,
+                    rb,
+                )
+                .map_err(to_pl_err)?;
+                Ok(out)
+            })
+        };
+
+        // handle non-Ruby threads
+        start_background_ruby_thread(&Ruby::get().unwrap());
+        let udf = ArcValue::new(rbfn.0);
+        let f = move |args: Args| {
+            if is_non_ruby_thread() {
+                let udf = udf.clone();
+                return run_in_ruby_thread(move || f(args, *udf.0));
+            }
+            f(args, *udf.0)
+        };
+
+        Self::new(f)
+    }
+}

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

@@ -0,0 +1,16 @@
+use magnus::{Module, RClass, Ruby, value::Lazy};
+
+static DATE: Lazy<RClass> = Lazy::new(|rb| rb.class_object().const_get("Date").unwrap());
+static DATETIME: Lazy<RClass> = Lazy::new(|rb| rb.class_object().const_get("DateTime").unwrap());
+
+pub(crate) fn bigdecimal(rb: &Ruby) -> Option<RClass> {
+    rb.class_object().const_get("BigDecimal").ok()
+}
+
+pub(crate) fn date(rb: &Ruby) -> RClass {
+    rb.get_inner(&DATE)
+}
+
+pub(crate) fn datetime(rb: &Ruby) -> RClass {
+    rb.get_inner(&DATETIME)
+}

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

@@ -0,0 +1,51 @@
+use std::any::Any;
+use std::ops::Deref;
+use std::sync::{Arc, LazyLock, RwLock};
+
+use magnus::{Value, value::Opaque};
+
+use crate::RbResult;
+
+pub type FromRuby = Arc<dyn Fn(Opaque<Value>) -> RbResult<Box<dyn Any>> + Send + Sync>;
+pub type ToRuby = Arc<dyn for<'a> Fn(&'a dyn Any) -> RbResult<Value> + Send + Sync>;
+
+#[derive(Clone)]
+pub struct FromRubyConvertRegistry {
+    #[allow(unused)]
+    pub file_provider_result: FromRuby,
+    pub series: FromRuby,
+    pub df: FromRuby,
+    pub dsl_plan: FromRuby,
+    pub schema: FromRuby,
+}
+
+#[derive(Clone)]
+pub struct ToRubyConvertRegistry {
+    pub df: ToRuby,
+    pub series: ToRuby,
+    pub dsl_plan: ToRuby,
+    pub schema: ToRuby,
+}
+
+#[derive(Clone)]
+pub struct RubyConvertRegistry {
+    pub from_rb: FromRubyConvertRegistry,
+    pub to_rb: ToRubyConvertRegistry,
+}
+
+static RUBY_CONVERT_REGISTRY: LazyLock<RwLock<Option<RubyConvertRegistry>>> =
+    LazyLock::new(Default::default);
+
+pub fn get_ruby_convert_registry() -> RubyConvertRegistry {
+    RUBY_CONVERT_REGISTRY
+        .deref()
+        .read()
+        .unwrap()
+        .as_ref()
+        .unwrap()
+        .clone()
+}
+
+pub fn register_converters(registry: RubyConvertRegistry) {
+    *RUBY_CONVERT_REGISTRY.deref().write().unwrap() = Some(registry);
+}

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

@@ -0,0 +1,11 @@
+use magnus::{Value, value::Opaque};
+
+pub struct RubyObject(pub Opaque<Value>);
+
+pub type RubyFunction = RubyObject;
+
+impl From<Value> for RubyObject {
+    fn from(value: Value) -> Self {
+        Self(value.into())
+    }
+}

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

@@ -0,0 +1,164 @@
+use std::sync::{Arc, OnceLock};
+
+use magnus::{Ruby, Value, value::Opaque, value::ReprValue};
+use polars::prelude::*;
+use polars_core::datatypes::{DataType, Field};
+use polars_core::frame::column::Column;
+use polars_core::schema::Schema;
+use polars_plan::dsl::udf::try_infer_udf_output_dtype;
+use polars_plan::prelude::*;
+use polars_utils::pl_str::PlSmallStr;
+
+use crate::ruby::gvl::GvlExt;
+use crate::ruby::utils::ArcValue;
+
+#[allow(clippy::type_complexity)]
+pub static mut CALL_COLUMNS_UDF_RUBY: Option<
+    fn(s: &[Column], output_dtype: Option<DataType>, lambda: Opaque<Value>) -> PolarsResult<Column>,
+> = None;
+
+#[allow(clippy::type_complexity)]
+pub static mut CALL_DF_UDF_RUBY: Option<
+    fn(s: DataFrame, lambda: Opaque<Value>) -> PolarsResult<DataFrame>,
+> = None;
+
+pub struct RubyUdfExpression {
+    ruby_function: ArcValue,
+    output_type: Option<DataTypeExpr>,
+    materialized_field: OnceLock<Field>,
+    is_elementwise: bool,
+    returns_scalar: bool,
+}
+
+impl RubyUdfExpression {
+    pub fn new(
+        lambda: Value,
+        output_type: Option<impl Into<DataTypeExpr>>,
+        is_elementwise: bool,
+        returns_scalar: bool,
+    ) -> Self {
+        let output_type = output_type.map(Into::into);
+        Self {
+            ruby_function: ArcValue::new(Opaque::from(lambda)),
+            output_type,
+            materialized_field: OnceLock::new(),
+            is_elementwise,
+            returns_scalar,
+        }
+    }
+}
+
+impl ColumnsUdf for RubyUdfExpression {
+    fn call_udf(&self, s: &mut [Column]) -> PolarsResult<Column> {
+        let func = unsafe { CALL_COLUMNS_UDF_RUBY.unwrap() };
+        let field = self
+            .materialized_field
+            .get()
+            .expect("should have been materialized at this point");
+        let mut out = func(
+            s,
+            self.materialized_field.get().map(|f| f.dtype.clone()),
+            *self.ruby_function.0,
+        )?;
+
+        let must_cast = out.dtype().matches_schema_type(field.dtype()).map_err(|_| {
+            polars_err!(
+                SchemaMismatch: "expected output type '{:?}', got '{:?}'; set `return_dtype` to the proper datatype",
+                field.dtype(), out.dtype(),
+            )
+        })?;
+        if must_cast {
+            out = out.cast(field.dtype())?;
+        }
+
+        Ok(out)
+    }
+}
+
+impl AnonymousColumnsUdf for RubyUdfExpression {
+    fn as_column_udf(self: Arc<Self>) -> Arc<dyn ColumnsUdf> {
+        self as _
+    }
+
+    fn deep_clone(self: Arc<Self>) -> Arc<dyn AnonymousColumnsUdf> {
+        Arc::new(Self {
+            ruby_function: ArcValue::new(Ruby::attach(|rb| {
+                // TODO fix
+                rb.get_inner(*self.ruby_function.0)
+                    .funcall::<_, _, Value>("dup", ())
+                    .map(Opaque::from)
+                    .unwrap()
+            })),
+            output_type: self.output_type.clone(),
+            materialized_field: OnceLock::new(),
+            is_elementwise: self.is_elementwise,
+            returns_scalar: self.returns_scalar,
+        }) as _
+    }
+
+    fn try_serialize(&self, _buf: &mut Vec<u8>) -> PolarsResult<()> {
+        todo!();
+    }
+
+    fn get_field(&self, input_schema: &Schema, fields: &[Field]) -> PolarsResult<Field> {
+        let field = match self.materialized_field.get() {
+            Some(f) => f.clone(),
+            None => {
+                let dtype = match self.output_type.as_ref() {
+                    None => {
+                        let func = unsafe { CALL_COLUMNS_UDF_RUBY.unwrap() };
+                        let f = |s: &[Column]| func(s, None, *self.ruby_function.0);
+                        try_infer_udf_output_dtype(&f as _, fields)?
+                    }
+                    Some(output_type) => output_type
+                        .clone()
+                        .into_datatype_with_self(input_schema, fields[0].dtype())?,
+                };
+
+                // Take the name of first field, just like `map_field`.
+                let name = fields[0].name();
+                let f = Field::new(name.clone(), dtype);
+                self.materialized_field.get_or_init(|| f.clone());
+                f
+            }
+        };
+        Ok(field)
+    }
+}
+
+pub trait RubyUdfExt {
+    #[allow(unused)]
+    fn map_ruby(self, func: RubyUdfExpression) -> Expr;
+
+    fn map_many_ruby(exprs: Vec<Expr>, func: RubyUdfExpression) -> Expr;
+}
+
+impl RubyUdfExt for Expr {
+    fn map_ruby(self, func: RubyUdfExpression) -> Expr {
+        Self::map_many_ruby(vec![self], func)
+    }
+
+    fn map_many_ruby(exprs: Vec<Expr>, func: RubyUdfExpression) -> Expr {
+        const NAME: &str = "ruby_udf";
+
+        let returns_scalar = func.returns_scalar;
+
+        let mut flags = FunctionFlags::default() | FunctionFlags::OPTIONAL_RE_ENTRANT;
+        if func.is_elementwise {
+            flags.set_elementwise();
+        }
+        if returns_scalar {
+            flags |= FunctionFlags::RETURNS_SCALAR;
+        }
+
+        Expr::AnonymousFunction {
+            input: exprs,
+            function: new_column_udf(func),
+            options: FunctionOptions {
+                flags,
+                ..Default::default()
+            },
+            fmt_str: Box::new(PlSmallStr::from(NAME)),
+        }
+    }
+}

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

@@ -0,0 +1,65 @@
+use std::any::Any;
+use std::sync::OnceLock;
+use std::sync::mpsc::{RecvTimeoutError, SyncSender, sync_channel};
+
+use magnus::Ruby;
+use magnus::error::RubyUnavailableError;
+
+use crate::ruby::gvl::GvlExt;
+
+type BackgroundMessage = (
+    Box<dyn FnOnce() -> Box<dyn Any + Send> + Send>,
+    SyncSender<Box<dyn Any + Send>>,
+);
+
+static BACKGROUND_THREAD_MAILBOX: OnceLock<SyncSender<BackgroundMessage>> = OnceLock::new();
+
+// TODO figure out better approach
+pub(crate) fn start_background_ruby_thread(rb: &Ruby) {
+    BACKGROUND_THREAD_MAILBOX.get_or_init(|| {
+        let (sender, receiver) = sync_channel::<BackgroundMessage>(0);
+
+        // TODO save reference to thread?
+        rb.thread_create_from_fn(move |rb2| {
+            rb2.detach(move || {
+                loop {
+                    match receiver.recv_timeout(std::time::Duration::from_millis(10)) {
+                        Ok((f, sender2)) => {
+                            sender2.send(f()).unwrap();
+                        }
+                        Err(RecvTimeoutError::Timeout) => {
+                            Ruby::attach(|rb3| rb3.thread_schedule());
+                        }
+                        Err(RecvTimeoutError::Disconnected) => {
+                            todo!();
+                        }
+                    }
+                }
+
+                #[allow(unreachable_code)]
+                ()
+            });
+        });
+
+        sender
+    });
+}
+
+pub(crate) fn run_in_ruby_thread<T, F>(func: F) -> T
+where
+    T: Send + 'static,
+    F: FnOnce() -> T + Send + 'static,
+{
+    let func2 = move || -> Box<dyn Any + Send> { Box::new(func()) };
+    let (sender, receiver) = sync_channel(0);
+    BACKGROUND_THREAD_MAILBOX
+        .get()
+        .unwrap()
+        .send((Box::new(func2), sender))
+        .unwrap();
+    *receiver.recv().unwrap().downcast().unwrap()
+}
+
+pub(crate) fn is_non_ruby_thread() -> bool {
+    matches!(Ruby::get(), Err(RubyUnavailableError::NonRubyThread))
+}

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

@@ -0,0 +1,39 @@
+use std::sync::Arc;
+
+use magnus::{Error, IntoValue, Ruby, Value, gc, value::Opaque};
+use polars::error::PolarsError;
+
+use crate::RbResult;
+use crate::ruby::gvl::GvlExt;
+
+pub(crate) fn to_pl_err(e: Error) -> PolarsError {
+    PolarsError::ComputeError(e.to_string().into())
+}
+
+#[derive(Clone)]
+pub struct ArcValue(pub Arc<Opaque<Value>>);
+
+impl ArcValue {
+    pub fn new(value: Opaque<Value>) -> Self {
+        let ob = Arc::new(value);
+        gc::register_address(&*ob);
+        Self(ob)
+    }
+}
+
+impl Drop for ArcValue {
+    fn drop(&mut self) {
+        // TODO use rb.gc_register_address
+        Ruby::attach(|_| gc::unregister_address(&*self.0));
+    }
+}
+
+pub trait TryIntoValue {
+    fn try_into_value_with(self, ruby: &Ruby) -> RbResult<Value>;
+}
+
+impl<T: IntoValue> TryIntoValue for T {
+    fn try_into_value_with(self, ruby: &Ruby) -> RbResult<Value> {
+        Ok(self.into_value_with(ruby))
+    }
+}