polars-df 0.21.0 → 0.22.0

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

@@ -1,4 +1,4 @@
-use magnus::{IntoValue, RArray, RHash, TryConvert, Value, r_hash::ForEach, typed_data::Obj};
+use magnus::{IntoValue, RArray, RHash, Ruby, TryConvert, Value, r_hash::ForEach, typed_data::Obj};
 use polars::io::{HiveOptions, RowIndex};
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
@@ -30,13 +30,12 @@ impl RbLazyFrame {
         source: Option<Value>,
         sources: Wrap<ScanSources>,
         infer_schema_length: Option<usize>,
-        batch_size: Option<Wrap<NonZeroUsize>>,
+        batch_size: Option<NonZeroUsize>,
         n_rows: Option<usize>,
         low_memory: bool,
         rechunk: bool,
         row_index: Option<(String, IdxSize)>,
     ) -> RbResult<Self> {
-        let batch_size = batch_size.map(|v| v.0);
         let row_index = row_index.map(|(name, offset)| RowIndex {
             name: name.into(),
             offset,
@@ -325,6 +324,30 @@ impl RbLazyFrame {
             .into())
     }
 
+    pub fn top_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
+        let ldf = self.ldf.borrow().clone();
+        let exprs = rb_exprs_to_exprs(by)?;
+        Ok(ldf
+            .top_k(
+                k,
+                exprs,
+                SortMultipleOptions::new().with_order_descending_multi(reverse),
+            )
+            .into())
+    }
+
+    pub fn bottom_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
+        let ldf = self.ldf.borrow().clone();
+        let exprs = rb_exprs_to_exprs(by)?;
+        Ok(ldf
+            .bottom_k(
+                k,
+                exprs,
+                SortMultipleOptions::new().with_order_descending_multi(reverse),
+            )
+            .into())
+    }
+
     pub fn cache(&self) -> Self {
         let ldf = self.ldf.borrow().clone();
         ldf.cache().into()
@@ -420,7 +443,7 @@ impl RbLazyFrame {
         let separator = u8::try_convert(arguments[3])?;
         let line_terminator = String::try_convert(arguments[4])?;
         let quote_char = u8::try_convert(arguments[5])?;
-        let batch_size = Wrap::<NonZeroUsize>::try_convert(arguments[6])?;
+        let batch_size = NonZeroUsize::try_convert(arguments[6])?;
         let datetime_format = Option::<String>::try_convert(arguments[7])?;
         let date_format = Option::<String>::try_convert(arguments[8])?;
         let time_format = Option::<String>::try_convert(arguments[9])?;
@@ -453,7 +476,7 @@ impl RbLazyFrame {
         let options = CsvWriterOptions {
             include_bom,
             include_header,
-            batch_size: batch_size.0,
+            batch_size,
             serialize_options,
         };
 
@@ -509,6 +532,11 @@ impl RbLazyFrame {
         ldf.filter(predicate.inner.clone()).into()
     }
 
+    pub fn remove(&self, predicate: &RbExpr) -> Self {
+        let ldf = self.ldf.borrow().clone();
+        ldf.remove(predicate.inner.clone()).into()
+    }
+
     pub fn select(&self, exprs: RArray) -> RbResult<Self> {
         let ldf = self.ldf.borrow().clone();
         let exprs = rb_exprs_to_exprs(exprs)?;
@@ -701,6 +729,20 @@ impl RbLazyFrame {
             .into())
     }
 
+    pub fn join_where(&self, other: &Self, predicates: RArray, suffix: String) -> RbResult<Self> {
+        let ldf = self.ldf.borrow().clone();
+        let other = other.ldf.borrow().clone();
+
+        let predicates = rb_exprs_to_exprs(predicates)?;
+
+        Ok(ldf
+            .join_builder()
+            .with(other)
+            .suffix(suffix)
+            .join_where(predicates)
+            .into())
+    }
+
     pub fn with_column(&self, expr: &RbExpr) -> Self {
         let ldf = self.ldf.borrow().clone();
         ldf.with_column(expr.inner.clone()).into()
@@ -816,6 +858,14 @@ impl RbLazyFrame {
         .into())
     }
 
+    pub fn drop_nans(&self, subset: Option<&RbSelector>) -> Self {
+        self.ldf
+            .borrow()
+            .clone()
+            .drop_nans(subset.map(|e| e.inner.clone()))
+            .into()
+    }
+
     pub fn drop_nulls(&self, subset: Option<&RbSelector>) -> Self {
         self.ldf
             .borrow()
@@ -880,19 +930,19 @@ impl RbLazyFrame {
         self.ldf.borrow().clone().into()
     }
 
-    pub fn collect_schema(&self) -> RbResult<RHash> {
-        let schema = self
+    pub fn collect_schema(ruby: &Ruby, rb_self: &Self) -> RbResult<RHash> {
+        let schema = rb_self
             .ldf
             .borrow_mut()
             .collect_schema()
             .map_err(RbPolarsErr::from)?;
 
-        let schema_dict = RHash::new();
+        let schema_dict = ruby.hash_new();
         schema.iter_fields().for_each(|fld| {
             schema_dict
                 .aset::<String, Value>(
                     fld.name().to_string(),
-                    Wrap(fld.dtype().clone()).into_value(),
+                    Wrap(fld.dtype().clone()).into_value_with(ruby),
                 )
                 .unwrap();
         });

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

@@ -1,8 +1,9 @@
 mod general;
+mod optflags;
 mod serde;
 mod sink;
 
-use polars::lazy::frame::LazyFrame;
+use polars::prelude::{LazyFrame, OptFlags};
 pub use sink::SinkTarget;
 use std::cell::RefCell;
 
@@ -12,6 +13,12 @@ pub struct RbLazyFrame {
     pub ldf: RefCell<LazyFrame>,
 }
 
+#[magnus::wrap(class = "Polars::RbOptFlags")]
+#[derive(Clone)]
+pub struct RbOptFlags {
+    pub inner: RefCell<OptFlags>,
+}
+
 impl From<LazyFrame> for RbLazyFrame {
     fn from(ldf: LazyFrame) -> Self {
         RbLazyFrame {
@@ -19,3 +26,11 @@ impl From<LazyFrame> for RbLazyFrame {
         }
     }
 }
+
+impl From<OptFlags> for RbOptFlags {
+    fn from(inner: OptFlags) -> Self {
+        RbOptFlags {
+            inner: RefCell::new(inner),
+        }
+    }
+}

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

@@ -0,0 +1,58 @@
+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.borrow_mut().remove(OptFlags::$flag);
+                })+
+            }
+
+            pub fn copy(&self) -> Self {
+                Self { inner: self.inner.clone() }
+            }
+
+            $(
+            pub fn $getter(&self) -> bool {
+                self.inner.borrow().contains(OptFlags::$flag)
+            }
+            pub fn $setter(&self, value: bool) {
+                self.inner.borrow_mut().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)
+    (COLLAPSE_JOINS, get_collapse_joins, set_collapse_joins, 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)
+
+    (EAGER, get_eager, set_eager, clear=true)
+    (NEW_STREAMING, get_streaming, set_streaming, clear=true)
+}

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

@@ -1,14 +1,38 @@
+use std::io::Read;
+#[cfg(feature = "serialize_binary")]
+use std::io::{BufReader, BufWriter};
+
 use magnus::Value;
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
-use std::io::Read;
 
 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
+            .borrow()
+            .logical_plan
+            .serialize_versioned(writer, Default::default())
+            .map_err(to_rb_err)
+    }
+
+    #[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();