polars-df 0.1.2 → 0.1.4

data/ext/polars/src/dataframe.rs

@@ -1,16 +1,14 @@
-use magnus::{r_hash::ForEach, Error, RArray, RHash, RString, Value};
+use magnus::{r_hash::ForEach, RArray, RHash, RString, Value};
 use polars::io::mmap::ReaderBytes;
 use polars::io::RowCount;
 use polars::prelude::*;
 use std::cell::RefCell;
-use std::fs::File;
-use std::io::{BufReader, BufWriter, Cursor};
+use std::io::{BufWriter, Cursor};
 use std::ops::Deref;
-use std::path::PathBuf;
 
 use crate::conversion::*;
 use crate::file::{get_file_like, get_mmap_bytes_reader};
-use crate::series::to_rbseries_collection;
+use crate::series::{to_rbseries_collection, to_series_collection};
 use crate::{series, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
 
 #[magnus::wrap(class = "Polars::RbDataFrame")]
@@ -137,13 +135,27 @@ impl RbDataFrame {
         Ok(df.into())
     }
 
-    pub fn read_parquet(path: PathBuf) -> RbResult<Self> {
-        let f = File::open(&path).map_err(|e| Error::runtime_error(e.to_string()))?;
-        let reader = BufReader::new(f);
-        ParquetReader::new(reader)
+    pub fn read_parquet(
+        rb_f: Value,
+        columns: Option<Vec<String>>,
+        projection: Option<Vec<usize>>,
+        n_rows: Option<usize>,
+        parallel: Wrap<ParallelStrategy>,
+        row_count: Option<(String, IdxSize)>,
+        low_memory: bool,
+    ) -> RbResult<Self> {
+        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+        let df = ParquetReader::new(mmap_bytes_r)
+            .with_projection(projection)
+            .with_columns(columns)
+            .read_parallel(parallel.0)
+            .with_n_rows(n_rows)
+            .with_row_count(row_count)
+            .set_low_memory(low_memory)
             .finish()
-            .map_err(RbPolarsErr::from)
-            .map(|v| v.into())
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
     }
 
     pub fn read_ipc(
@@ -313,6 +325,55 @@ impl RbDataFrame {
         Ok(())
     }
 
+    pub fn row_tuple(&self, idx: i64) -> Value {
+        let idx = if idx < 0 {
+            (self.df.borrow().height() as i64 + idx) as usize
+        } else {
+            idx as usize
+        };
+        RArray::from_vec(
+            self.df
+                .borrow()
+                .get_columns()
+                .iter()
+                .map(|s| match s.dtype() {
+                    DataType::Object(_) => {
+                        let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
+                        obj.unwrap().to_object()
+                    }
+                    _ => Wrap(s.get(idx)).into(),
+                })
+                .collect(),
+        )
+        .into()
+    }
+
+    pub fn row_tuples(&self) -> Value {
+        let df = &self.df;
+        RArray::from_vec(
+            (0..df.borrow().height())
+                .map(|idx| {
+                    RArray::from_vec(
+                        self.df
+                            .borrow()
+                            .get_columns()
+                            .iter()
+                            .map(|s| match s.dtype() {
+                                DataType::Object(_) => {
+                                    let obj: Option<&ObjectValue> =
+                                        s.get_object(idx).map(|any| any.into());
+                                    obj.unwrap().to_object()
+                                }
+                                _ => Wrap(s.get(idx)).into(),
+                            })
+                            .collect(),
+                    )
+                })
+                .collect(),
+        )
+        .into()
+    }
+
     pub fn write_parquet(
         &self,
         rb_f: Value,
@@ -338,6 +399,86 @@ impl RbDataFrame {
         Ok(())
     }
 
+    pub fn add(&self, s: &RbSeries) -> RbResult<Self> {
+        let df = (&*self.df.borrow() + &*s.series.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn sub(&self, s: &RbSeries) -> RbResult<Self> {
+        let df = (&*self.df.borrow() - &*s.series.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn div(&self, s: &RbSeries) -> RbResult<Self> {
+        let df = (&*self.df.borrow() / &*s.series.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn mul(&self, s: &RbSeries) -> RbResult<Self> {
+        let df = (&*self.df.borrow() * &*s.series.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn rem(&self, s: &RbSeries) -> RbResult<Self> {
+        let df = (&*self.df.borrow() % &*s.series.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn add_df(&self, s: &Self) -> RbResult<Self> {
+        let df = (&*self.df.borrow() + &*s.df.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn sub_df(&self, s: &Self) -> RbResult<Self> {
+        let df = (&*self.df.borrow() - &*s.df.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn div_df(&self, s: &Self) -> RbResult<Self> {
+        let df = (&*self.df.borrow() / &*s.df.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn mul_df(&self, s: &Self) -> RbResult<Self> {
+        let df = (&*self.df.borrow() * &*s.df.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn rem_df(&self, s: &Self) -> RbResult<Self> {
+        let df = (&*self.df.borrow() % &*s.df.borrow()).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn sample_n(
+        &self,
+        n: usize,
+        with_replacement: bool,
+        shuffle: bool,
+        seed: Option<u64>,
+    ) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .sample_n(n, with_replacement, shuffle, seed)
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn sample_frac(
+        &self,
+        frac: f64,
+        with_replacement: bool,
+        shuffle: bool,
+        seed: Option<u64>,
+    ) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .sample_frac(frac, with_replacement, shuffle, seed)
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
     pub fn rechunk(&self) -> Self {
         self.df.borrow().agg_chunks().into()
     }
@@ -393,6 +534,73 @@ impl RbDataFrame {
         self.df.borrow().width()
     }
 
+    pub fn hstack_mut(&self, columns: RArray) -> RbResult<()> {
+        let columns = to_series_collection(columns)?;
+        self.df
+            .borrow_mut()
+            .hstack_mut(&columns)
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn hstack(&self, columns: RArray) -> RbResult<Self> {
+        let columns = to_series_collection(columns)?;
+        let df = self
+            .df
+            .borrow()
+            .hstack(&columns)
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn extend(&self, df: &RbDataFrame) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .extend(&df.df.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn vstack_mut(&self, df: &RbDataFrame) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .vstack_mut(&df.df.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn vstack(&self, df: &RbDataFrame) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .vstack(&df.df.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn drop_in_place(&self, name: String) -> RbResult<RbSeries> {
+        let s = self
+            .df
+            .borrow_mut()
+            .drop_in_place(&name)
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbSeries::new(s))
+    }
+
+    pub fn drop_nulls(&self, subset: Option<Vec<String>>) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .drop_nulls(subset.as_ref().map(|s| s.as_ref()))
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn drop(&self, name: String) -> RbResult<Self> {
+        let df = self.df.borrow().drop(&name).map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
     pub fn select_at_idx(&self, idx: usize) -> Option<RbSeries> {
         self.df
             .borrow()
@@ -400,6 +608,10 @@ impl RbDataFrame {
             .map(|s| RbSeries::new(s.clone()))
     }
 
+    pub fn find_idx_by_name(&self, name: String) -> Option<usize> {
+        self.df.borrow().find_idx_by_name(&name)
+    }
+
     // TODO remove clone
     pub fn column(&self, name: String) -> RbResult<RbSeries> {
         self.df
@@ -702,6 +914,11 @@ impl RbDataFrame {
         Ok(out.into())
     }
 
+    pub fn to_struct(&self, name: String) -> RbSeries {
+        let s = self.df.borrow().clone().into_struct(&name);
+        s.into_series().into()
+    }
+
     pub fn unnest(&self, names: Vec<String>) -> RbResult<Self> {
         let df = self.df.borrow().unnest(names).map_err(RbPolarsErr::from)?;
         Ok(df.into())

data/ext/polars/src/error.rs

@@ -22,6 +22,10 @@ impl RbPolarsErr {
     pub fn other(message: String) -> Error {
         Error::runtime_error(message)
     }
+
+    pub fn todo() -> Error {
+        Error::runtime_error("not implemented yet")
+    }
 }
 
 pub struct RbValueError {}

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

@@ -211,7 +211,7 @@ impl RbLazyFrame {
         projection_pushdown: bool,
         simplify_expr: bool,
         slice_pushdown: bool,
-        _cse: bool,
+        cse: bool,
         allow_streaming: bool,
     ) -> RbLazyFrame {
         let ldf = self.ldf.clone();
@@ -220,7 +220,7 @@ impl RbLazyFrame {
             .with_predicate_pushdown(predicate_pushdown)
             .with_simplify_expr(simplify_expr)
             .with_slice_pushdown(slice_pushdown)
-            // .with_common_subplan_elimination(cse)
+            .with_common_subplan_elimination(cse)
             .with_streaming(allow_streaming)
             .with_projection_pushdown(projection_pushdown);
         ldf.into()
@@ -520,9 +520,9 @@ impl RbLazyFrame {
         Ok(self.get_schema()?.iter_names().cloned().collect())
     }
 
-    pub fn dtypes(&self) -> RbResult<Vec<String>> {
+    pub fn dtypes(&self) -> RbResult<Vec<Value>> {
         let schema = self.get_schema()?;
-        let iter = schema.iter_dtypes().map(|dt| dt.to_string());
+        let iter = schema.iter_dtypes().map(|dt| Wrap(dt.clone()).into());
         Ok(iter.collect())
     }
 
@@ -533,7 +533,7 @@ impl RbLazyFrame {
         schema.iter_fields().for_each(|fld| {
             // TODO remove unwrap
             schema_dict
-                .aset(fld.name().clone(), fld.data_type().to_string())
+                .aset::<String, Value>(fld.name().clone(), Wrap(fld.data_type().clone()).into())
                 .unwrap();
         });
         Ok(schema_dict)

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

@@ -1,4 +1,4 @@
-use magnus::{RArray, RString, Value};
+use magnus::{class, RArray, RString, Value};
 use polars::chunked_array::ops::SortOptions;
 use polars::lazy::dsl;
 use polars::lazy::dsl::Operator;
@@ -8,7 +8,8 @@ use polars::series::ops::NullBehavior;
 use crate::conversion::*;
 use crate::lazy::apply::*;
 use crate::lazy::utils::rb_exprs_to_exprs;
-use crate::RbResult;
+use crate::utils::reinterpret;
+use crate::{RbResult, RbSeries};
 
 #[magnus::wrap(class = "Polars::RbExpr")]
 #[derive(Clone)]
@@ -658,6 +659,65 @@ impl RbExpr {
         self.inner.clone().str().starts_with(sub).into()
     }
 
+    pub fn str_hex_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8().map(|s| s.hex_encode().into_series()),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.hex_encode")
+            .into()
+    }
+
+    pub fn str_hex_decode(&self, strict: Option<bool>) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8()?.hex_decode(strict).map(|s| s.into_series()),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.hex_decode")
+            .into()
+    }
+
+    pub fn str_base64_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8().map(|s| s.base64_encode().into_series()),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.base64_encode")
+            .into()
+    }
+
+    pub fn str_base64_decode(&self, strict: Option<bool>) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8()?.base64_decode(strict).map(|s| s.into_series()),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.base64_decode")
+            .into()
+    }
+
+    pub fn str_json_path_match(&self, pat: String) -> Self {
+        let function = move |s: Series| {
+            let ca = s.utf8()?;
+            match ca.json_path_match(&pat) {
+                Ok(ca) => Ok(ca.into_series()),
+                Err(e) => Err(PolarsError::ComputeError(format!("{:?}", e).into())),
+            }
+        };
+        self.clone()
+            .inner
+            .map(function, GetOutput::from_type(DataType::Utf8))
+            .with_fmt("str.json_path_match")
+            .into()
+    }
+
     pub fn str_extract(&self, pat: String, group_index: usize) -> Self {
         self.inner.clone().str().extract(&pat, group_index).into()
     }
@@ -886,6 +946,23 @@ impl RbExpr {
         self.inner.clone().dt().round(&every, &offset).into()
     }
 
+    pub fn dot(&self, other: &RbExpr) -> Self {
+        self.inner.clone().dot(other.inner.clone()).into()
+    }
+
+    pub fn reinterpret(&self, signed: bool) -> Self {
+        let function = move |s: Series| reinterpret(&s, signed);
+        let dt = if signed {
+            DataType::Int64
+        } else {
+            DataType::UInt64
+        };
+        self.clone()
+            .inner
+            .map(function, GetOutput::from_type(dt))
+            .into()
+    }
+
     pub fn mode(&self) -> Self {
         self.inner.clone().mode().into()
     }
@@ -1194,6 +1271,28 @@ impl RbExpr {
             .into()
     }
 
+    pub fn lst_to_struct(
+        &self,
+        width_strat: Wrap<ListToStructWidthStrategy>,
+        _name_gen: Option<Value>,
+    ) -> RbResult<Self> {
+        // TODO fix
+        let name_gen = None;
+        // let name_gen = name_gen.map(|lambda| {
+        //     Arc::new(move |idx: usize| {
+        //         let out: Value = lambda.funcall("call", (idx,)).unwrap();
+        //         out.try_convert::<String>().unwrap()
+        //     }) as NameGenerator
+        // });
+
+        Ok(self
+            .inner
+            .clone()
+            .arr()
+            .to_struct(width_strat.0, name_gen)
+            .into())
+    }
+
     pub fn rank(&self, method: Wrap<RankMethod>, reverse: bool) -> Self {
         let options = RankOptions {
             method: method.0,
@@ -1305,6 +1404,21 @@ impl RbExpr {
         self.inner.clone().ewm_var(options).into()
     }
 
+    pub fn extend_constant(&self, value: Wrap<AnyValue>, n: usize) -> Self {
+        let value = Value::from(value);
+        self.inner
+            .clone()
+            .apply(
+                move |s| {
+                    let value = value.try_convert::<Wrap<AnyValue>>().unwrap().0;
+                    s.extend_constant(value, n)
+                },
+                GetOutput::same_type(),
+            )
+            .with_fmt("extend")
+            .into()
+    }
+
     pub fn any(&self) -> Self {
         self.inner.clone().any().into()
     }
@@ -1332,6 +1446,10 @@ impl RbExpr {
     pub fn exp(&self) -> Self {
         self.inner.clone().exp().into()
     }
+
+    pub fn entropy(&self, base: f64, normalize: bool) -> Self {
+        self.inner.clone().entropy(base, normalize).into()
+    }
 }
 
 pub fn col(name: String) -> RbExpr {
@@ -1365,8 +1483,24 @@ pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
 pub fn lit(value: Value) -> RbResult<RbExpr> {
     if value.is_nil() {
         Ok(dsl::lit(Null {}).into())
+    } else if let Ok(series) = value.try_convert::<&RbSeries>() {
+        Ok(dsl::lit(series.series.borrow().clone()).into())
     } else if let Some(v) = RString::from_value(value) {
         Ok(dsl::lit(v.try_convert::<String>()?).into())
+    } else if value.is_kind_of(class::integer()) {
+        match value.try_convert::<i64>() {
+            Ok(val) => {
+                if val > 0 && val < i32::MAX as i64 || val < 0 && val > i32::MIN as i64 {
+                    Ok(dsl::lit(val as i32).into())
+                } else {
+                    Ok(dsl::lit(val).into())
+                }
+            }
+            _ => {
+                let val = value.try_convert::<u64>()?;
+                Ok(dsl::lit(val).into())
+            }
+        }
     } else {
         Ok(dsl::lit(value.try_convert::<f64>()?).into())
     }
@@ -1376,6 +1510,27 @@ pub fn arange(low: &RbExpr, high: &RbExpr, step: usize) -> RbExpr {
     polars::lazy::dsl::arange(low.inner.clone(), high.inner.clone(), step).into()
 }
 
+pub fn repeat(value: Value, n_times: &RbExpr) -> RbResult<RbExpr> {
+    if value.is_nil() {
+        Ok(polars::lazy::dsl::repeat(Null {}, n_times.inner.clone()).into())
+    } else {
+        todo!();
+    }
+}
+
+pub fn pearson_corr(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
+    polars::lazy::dsl::pearson_corr(a.inner.clone(), b.inner.clone(), ddof).into()
+}
+
+pub fn spearman_rank_corr(a: &RbExpr, b: &RbExpr, ddof: u8, propagate_nans: bool) -> RbExpr {
+    polars::lazy::dsl::spearman_rank_corr(a.inner.clone(), b.inner.clone(), ddof, propagate_nans)
+        .into()
+}
+
+pub fn cov(a: &RbExpr, b: &RbExpr) -> RbExpr {
+    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone()).into()
+}
+
 #[magnus::wrap(class = "Polars::RbWhen")]
 #[derive(Clone)]
 pub struct RbWhen {