polars-df 0.15.0 → 0.26.0

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

@@ -1,19 +1,72 @@
+mod exitable;
 mod general;
+mod optflags;
 mod serde;
+mod sink;
 
-use polars::lazy::frame::LazyFrame;
-use std::cell::RefCell;
+pub use exitable::RbInProcessQuery;
+pub use general::RbCollectBatches;
+use magnus::{TryConvert, Value};
+use parking_lot::RwLock;
+use polars::prelude::{Engine, LazyFrame, OptFlags};
+
+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>,
+}
+
+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()),
+        }
+    }
+}
+
+impl From<OptFlags> for RbOptFlags {
+    fn from(inner: OptFlags) -> Self {
+        RbOptFlags {
+            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

@@ -0,0 +1,59 @@
+use parking_lot::RwLock;
+use polars::prelude::OptFlags;
+
+use super::RbOptFlags;
+
+macro_rules! flag_getter_setters {
+    ($(($flag:ident, $getter:ident, $setter:ident, clear=$clear:literal))+) => {
+        impl RbOptFlags {
+            pub fn empty() -> Self {
+                Self {
+                    inner: OptFlags::empty().into()
+                }
+            }
+
+            #[allow(clippy::should_implement_trait)]
+            pub fn default() -> Self {
+                Self { inner: OptFlags::default().into() }
+            }
+
+            pub fn no_optimizations(&self) {
+                $(if $clear {
+                    self.inner.write().remove(OptFlags::$flag);
+                })+
+            }
+
+            pub fn copy(&self) -> Self {
+                Self { inner: RwLock::new(self.inner.read().clone()) }
+            }
+
+            $(
+            pub fn $getter(&self) -> bool {
+                self.inner.read().contains(OptFlags::$flag)
+            }
+            pub fn $setter(&self, value: bool) {
+                self.inner.write().set(OptFlags::$flag, value)
+            }
+            )+
+        }
+    };
+}
+
+flag_getter_setters! {
+    (TYPE_COERCION, get_type_coercion, set_type_coercion, clear=false)
+    (TYPE_CHECK, get_type_check, set_type_check, clear=false)
+
+    (PROJECTION_PUSHDOWN, get_projection_pushdown, set_projection_pushdown, clear=true)
+    (PREDICATE_PUSHDOWN, get_predicate_pushdown, set_predicate_pushdown, clear=true)
+    (CLUSTER_WITH_COLUMNS, get_cluster_with_columns, set_cluster_with_columns, clear=true)
+    (SIMPLIFY_EXPR, get_simplify_expression, set_simplify_expression, clear=true)
+    (SLICE_PUSHDOWN, get_slice_pushdown, set_slice_pushdown, clear=true)
+    (COMM_SUBPLAN_ELIM, get_comm_subplan_elim, set_comm_subplan_elim, clear=true)
+    (COMM_SUBEXPR_ELIM, get_comm_subexpr_elim, set_comm_subexpr_elim, clear=true)
+    (CHECK_ORDER_OBSERVE, get_check_order_observe, set_check_order_observe, clear=true)
+    (FAST_PROJECTION, get_fast_projection, set_fast_projection, clear=true)
+    (SORT_COLLAPSE, get_sort_collapse, set_sort_collapse, clear=true)
+
+    (EAGER, get_eager, set_eager, clear=true)
+    (STREAMING, get_streaming, set_streaming, clear=true)
+}

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

@@ -1,14 +1,46 @@
+#[cfg(feature = "serialize_binary")]
+use std::io::BufReader;
+use std::io::{BufWriter, Read};
+
 use magnus::Value;
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
-use std::io::Read;
 
+use crate::exceptions::ComputeError;
 use crate::file::get_file_like;
+#[cfg(feature = "serialize_binary")]
+use crate::utils::to_rb_err;
 use crate::{RbLazyFrame, RbResult, RbValueError};
 
 impl RbLazyFrame {
-    // TODO change to serialize_json
-    pub fn read_json(rb_f: Value) -> RbResult<Self> {
+    #[cfg(feature = "serialize_binary")]
+    pub fn serialize_binary(&self, rb_f: Value) -> RbResult<()> {
+        let file = get_file_like(rb_f, true)?;
+        let writer = BufWriter::new(file);
+        self.ldf
+            .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)?;
+        let reader = BufReader::new(file);
+
+        let lp: DslPlan = DslPlan::deserialize_versioned(reader).map_err(to_rb_err)?;
+        Ok(LazyFrame::from(lp).into())
+    }
+
+    pub fn deserialize_json(rb_f: Value) -> RbResult<Self> {
         // it is faster to first read to memory and then parse: https://github.com/serde-rs/json/issues/160
         // so don't bother with files.
         let mut json = String::new();
@@ -25,7 +57,7 @@ impl RbLazyFrame {
         let json = unsafe { std::mem::transmute::<&'_ str, &'static str>(json.as_str()) };
 
         let lp = serde_json::from_str::<DslPlan>(json)
-            .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
+            .map_err(|err| RbValueError::new_err(format!("{err:?}")))?;
         Ok(LazyFrame::from(lp).into())
     }
 }

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

@@ -0,0 +1,46 @@
+use std::sync::{Arc, Mutex};
+
+use magnus::{Ruby, TryConvert, Value};
+use polars::prelude::sync_on_close::SyncOnCloseType;
+use polars::prelude::{PlRefPath, SpecialEq};
+
+use crate::prelude::Wrap;
+use crate::ruby::gvl::GvlExt;
+use crate::{RbResult, RbValueError};
+
+impl TryConvert for Wrap<polars_plan::dsl::SinkTarget> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        if let Ok(v) = String::try_convert(ob) {
+            Ok(Wrap(polars::prelude::SinkTarget::Path(PlRefPath::new(&v))))
+        } else {
+            let writer = Ruby::attach(|rb| {
+                let rb_f = ob;
+                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))),
+            ))))
+        }
+    }
+}
+
+impl TryConvert for Wrap<SyncOnCloseType> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match String::try_convert(ob)?.as_str() {
+            "none" => SyncOnCloseType::None,
+            "data" => SyncOnCloseType::Data,
+            "all" => SyncOnCloseType::All,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "`sync_on_close` must be one of {{'none', 'data', 'all'}}, got {v}",
+                )));
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}

data/ext/polars/src/lazygroupby.rs

@@ -1,29 +1,58 @@
-use magnus::RArray;
-use polars::lazy::frame::LazyGroupBy;
-use std::cell::RefCell;
+use std::sync::Arc;
 
-use crate::expr::rb_exprs_to_exprs;
+use magnus::{RArray, Value};
+use polars::lazy::frame::{LazyFrame, LazyGroupBy};
+use polars::prelude::{PlanCallback, Schema};
+
+use crate::conversion::Wrap;
+use crate::error::RbPolarsErr;
+use crate::expr::ToExprs;
+use crate::ruby::plan_callback::PlanCallbackExt;
+use crate::ruby::ruby_function::RubyObject;
 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 having(&self, predicates: RArray) -> RbResult<RbLazyGroupBy> {
+        let mut lgb = self.lgb.clone().unwrap();
+        let predicates = predicates.to_exprs()?;
+        for predicate in predicates.into_iter() {
+            lgb = lgb.having(predicate);
+        }
+        Ok(RbLazyGroupBy { lgb: Some(lgb) })
+    }
+
     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()
     }
+
+    pub fn map_groups(&self, lambda: Value, schema: Option<Wrap<Schema>>) -> RbResult<RbLazyFrame> {
+        let lgb = self.lgb.clone().unwrap();
+        let schema = match schema {
+            Some(schema) => Arc::new(schema.0),
+            None => LazyFrame::from(lgb.logical_plan.clone())
+                .collect_schema()
+                .map_err(RbPolarsErr::from)?,
+        };
+
+        let function = RubyObject::from(lambda);
+
+        Ok(lgb.apply(PlanCallback::new_ruby(function), schema).into())
+    }
 }