polars-df 0.23.0 → 0.24.0

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

@@ -1,28 +1,55 @@
+mod exitable;
 mod general;
 mod optflags;
 mod serde;
 mod sink;
 
-use polars::prelude::{LazyFrame, OptFlags};
+pub use exitable::RbInProcessQuery;
+use magnus::{TryConvert, Value};
+use parking_lot::RwLock;
+use polars::prelude::{Engine, LazyFrame, OptFlags};
 pub use sink::SinkTarget;
-use std::cell::RefCell;
+
+use crate::prelude::Wrap;
+use crate::{RbResult, RbValueError};
 
 #[magnus::wrap(class = "Polars::RbLazyFrame")]
-#[derive(Clone)]
+#[repr(transparent)]
 pub struct RbLazyFrame {
-    pub ldf: RefCell<LazyFrame>,
+    pub ldf: RwLock<LazyFrame>,
 }
 
-#[magnus::wrap(class = "Polars::RbOptFlags")]
-#[derive(Clone)]
-pub struct RbOptFlags {
-    pub inner: RefCell<OptFlags>,
+impl Clone for RbLazyFrame {
+    fn clone(&self) -> Self {
+        Self {
+            ldf: RwLock::new(self.ldf.read().clone()),
+        }
+    }
 }
 
 impl From<LazyFrame> for RbLazyFrame {
     fn from(ldf: LazyFrame) -> Self {
         RbLazyFrame {
-            ldf: RefCell::new(ldf),
+            ldf: RwLock::new(ldf),
+        }
+    }
+}
+
+impl From<RbLazyFrame> for LazyFrame {
+    fn from(pldf: RbLazyFrame) -> Self {
+        pldf.ldf.into_inner()
+    }
+}
+
+#[magnus::wrap(class = "Polars::RbOptFlags")]
+pub struct RbOptFlags {
+    pub inner: RwLock<OptFlags>,
+}
+
+impl Clone for RbOptFlags {
+    fn clone(&self) -> Self {
+        Self {
+            inner: RwLock::new(*self.inner.read()),
         }
     }
 }
@@ -30,7 +57,16 @@ impl From<LazyFrame> for RbLazyFrame {
 impl From<OptFlags> for RbOptFlags {
     fn from(inner: OptFlags) -> Self {
         RbOptFlags {
-            inner: RefCell::new(inner),
+            inner: RwLock::new(inner),
         }
     }
 }
+
+impl TryConvert for Wrap<Engine> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = String::try_convert(ob)?
+            .parse()
+            .map_err(RbValueError::new_err)?;
+        Ok(Wrap(parsed))
+    }
+}

data/ext/polars/src/lazyframe/optflags.rs

@@ -1,3 +1,4 @@
+use parking_lot::RwLock;
 use polars::prelude::OptFlags;
 
 use super::RbOptFlags;
@@ -18,20 +19,20 @@ macro_rules! flag_getter_setters {
 
             pub fn no_optimizations(&self) {
                 $(if $clear {
-                    self.inner.borrow_mut().remove(OptFlags::$flag);
+                    self.inner.write().remove(OptFlags::$flag);
                 })+
             }
 
             pub fn copy(&self) -> Self {
-                Self { inner: self.inner.clone() }
+                Self { inner: RwLock::new(self.inner.read().clone()) }
             }
 
             $(
             pub fn $getter(&self) -> bool {
-                self.inner.borrow().contains(OptFlags::$flag)
+                self.inner.read().contains(OptFlags::$flag)
             }
             pub fn $setter(&self, value: bool) {
-                self.inner.borrow_mut().set(OptFlags::$flag, value)
+                self.inner.write().set(OptFlags::$flag, value)
             }
             )+
         }

data/ext/polars/src/lazyframe/serde.rs

@@ -1,11 +1,12 @@
-use std::io::Read;
 #[cfg(feature = "serialize_binary")]
-use std::io::{BufReader, BufWriter};
+use std::io::BufReader;
+use std::io::{BufWriter, Read};
 
 use magnus::Value;
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
 
+use crate::exceptions::ComputeError;
 use crate::file::get_file_like;
 #[cfg(feature = "serialize_binary")]
 use crate::utils::to_rb_err;
@@ -17,12 +18,19 @@ impl RbLazyFrame {
         let file = get_file_like(rb_f, true)?;
         let writer = BufWriter::new(file);
         self.ldf
-            .borrow()
+            .read()
             .logical_plan
             .serialize_versioned(writer, Default::default())
             .map_err(to_rb_err)
     }
 
+    pub fn serialize_json(&self, rb_f: Value) -> RbResult<()> {
+        let file = get_file_like(rb_f, true)?;
+        let writer = BufWriter::new(file);
+        serde_json::to_writer(writer, &self.ldf.read().logical_plan)
+            .map_err(|err| ComputeError::new_err(err.to_string()))
+    }
+
     #[cfg(feature = "serialize_binary")]
     pub fn deserialize_binary(rb_f: Value) -> RbResult<Self> {
         let file = get_file_like(rb_f, false)?;

data/ext/polars/src/lazyframe/sink.rs

@@ -1,11 +1,12 @@
 use std::sync::{Arc, Mutex};
 
-use magnus::{RHash, TryConvert, Value};
+use magnus::{RHash, Ruby, TryConvert, Value};
 use polars::prelude::sync_on_close::SyncOnCloseType;
 use polars::prelude::{PlPath, SinkOptions, SpecialEq};
 use polars_utils::plpath::PlPathRef;
 
 use crate::prelude::Wrap;
+use crate::utils::RubyAttach;
 use crate::{RbResult, RbValueError};
 
 #[derive(Clone)]
@@ -18,10 +19,14 @@ impl TryConvert for Wrap<polars_plan::dsl::SinkTarget> {
         if let Ok(v) = String::try_convert(ob) {
             Ok(Wrap(polars::prelude::SinkTarget::Path(PlPath::new(&v))))
         } else {
-            let writer = {
+            let writer = Ruby::attach(|rb| {
                 let rb_f = ob;
-                RbResult::Ok(crate::file::try_get_rbfile(rb_f, true)?.0.into_writeable())
-            }?;
+                RbResult::Ok(
+                    crate::file::try_get_rbfile(rb, rb_f, true)?
+                        .0
+                        .into_writeable(),
+                )
+            })?;
 
             Ok(Wrap(polars_plan::prelude::SinkTarget::Dyn(SpecialEq::new(
                 Arc::new(Mutex::new(Some(writer))),
@@ -71,7 +76,7 @@ impl TryConvert for Wrap<SinkOptions> {
 
         if parsed.len() != 3 {
             return Err(RbValueError::new_err(
-                "`sink_options` must be a dictionary with the exactly 3 field.",
+                "`sink_options` must be a hash with the exactly 3 field.",
             ));
         }
 

data/ext/polars/src/lazygroupby.rs

@@ -1,29 +1,28 @@
 use magnus::RArray;
 use polars::lazy::frame::LazyGroupBy;
-use std::cell::RefCell;
 
-use crate::expr::rb_exprs_to_exprs;
+use crate::expr::ToExprs;
 use crate::{RbLazyFrame, RbResult};
 
 #[magnus::wrap(class = "Polars::RbLazyGroupBy")]
 pub struct RbLazyGroupBy {
-    pub lgb: RefCell<Option<LazyGroupBy>>,
+    pub lgb: Option<LazyGroupBy>,
 }
 
 impl RbLazyGroupBy {
     pub fn agg(&self, aggs: RArray) -> RbResult<RbLazyFrame> {
-        let lgb = self.lgb.borrow_mut().take().unwrap();
-        let aggs = rb_exprs_to_exprs(aggs)?;
+        let lgb = self.lgb.clone().unwrap();
+        let aggs = aggs.to_exprs()?;
         Ok(lgb.agg(aggs).into())
     }
 
     pub fn head(&self, n: usize) -> RbLazyFrame {
-        let lgb = self.lgb.take().unwrap();
+        let lgb = self.lgb.clone().unwrap();
         lgb.head(Some(n)).into()
     }
 
     pub fn tail(&self, n: usize) -> RbLazyFrame {
-        let lgb = self.lgb.take().unwrap();
+        let lgb = self.lgb.clone().unwrap();
         lgb.tail(Some(n)).into()
     }
 }