polars-df 0.1.2 → 0.1.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 510c7a761553fb3a49919add842d520da1518c5df6dd37f93af338e928c7a207
-  data.tar.gz: 3e724d3c2553bb6b4a587f056a04977990b85103ea9ae231166746407e1c0b1e
+  metadata.gz: 9e6fb732e5dafe2fde285322554bd9159483cbbdf17d6e2bba9cba9a83563b47
+  data.tar.gz: 1b4249d0c0100f136973c601b8404cb6d92abc632d5ed0476bd93bc5360a11dc
 SHA512:
-  metadata.gz: 59391ecd7a3d14372ad24693240ebeae4d83c59d55b37876fc6f216dd576897113607f7b5f6f6a3d7a5d75f96def72b1f2659f3de71e6a975a2753719879fa33
-  data.tar.gz: 5a2e8cdf3aed01f16823be62067b0ec01e4e0dc3bfa4dc27929dfbe42d7c92e62f6e7a09bc78533d220a3bfc9890f9f0fb069482a338fbd0b55603aa42fddc57
+  metadata.gz: d9414d6f60c489e2b3b72885288822083ba8c04bac4053f4e34c1d53ee805d164f17fe4b8b3a8f4ff562550bcc657f374bea6e250b52985367f601ea50e3037f
+  data.tar.gz: 9e3a7cfe105f03ec20e9c26aa38c1475074ccc1ea057a170a97b7068b41943d561d50af49bb1d1f74b7705809dc1375900f542ab93683ba627dea080274f6d91

data/.yardopts

@@ -0,0 +1,3 @@
+--no-private
+--markup markdown
+--embed-mixins

data/CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.1.3 (2022-11-27)
+
+- Added more methods
+
 ## 0.1.2 (2022-11-25)
 
 - Added more methods

data/Cargo.lock

@@ -1160,7 +1160,7 @@ dependencies = [
 
 [[package]]
 name = "polars"
-version = "0.1.2"
+version = "0.1.3"
 dependencies = [
  "magnus",
  "polars 0.25.1",
@@ -1217,6 +1217,7 @@ dependencies = [
  "rayon",
  "regex",
  "serde",
+ "serde_json",
  "smartstring",
  "thiserror",
 ]

data/README.md

@@ -27,7 +27,7 @@ Polars.read_csv("iris.csv")
   .collect
 ```
 
-You can follow [Polars tutorials](https://pola-rs.github.io/polars-book/user-guide/introduction.html) and convert the code to Ruby in many cases. Feel free to open an issue if you run into problems. Note that many methods and options are missing at the moment.
+You can follow [Polars tutorials](https://pola-rs.github.io/polars-book/user-guide/introduction.html) and convert the code to Ruby in many cases. Feel free to open an issue if you run into problems. Some methods are missing at the moment.
 
 ## Examples
 

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.1.2"
+version = "0.1.3"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
 publish = false
@@ -17,10 +17,12 @@ version = "0.25.1"
 features = [
     "abs",
     "arange",
+    "arg_where",
     "concat_str",
     "csv-file",
     "cum_agg",
     "cumulative_eval",
+    "dataframe_arithmetic",
     "date_offset",
     "diagonal_concat",
     "diff",
@@ -38,16 +40,20 @@ features = [
     "lazy",
     "lazy_regex",
     "list_eval",
+    "list_to_struct",
     "log",
     "meta",
     "mode",
     "moment",
+    "object",
     "parquet",
     "partition_by",
     "pct_change",
     "product",
+    "propagate_nans",
     "random",
     "rank",
+    "reinterpret",
     "repeat_by",
     "rolling_window",
     "round_series",

data/ext/polars/src/conversion.rs

@@ -1,4 +1,5 @@
-use magnus::{RArray, Symbol, TryConvert, Value, QNIL};
+use magnus::{class, RArray, Symbol, TryConvert, Value, QNIL};
+use polars::chunked_array::object::PolarsObjectSafe;
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
 use polars::datatypes::AnyValue;
 use polars::frame::DataFrame;
@@ -79,6 +80,13 @@ impl From<Wrap<AnyValue<'_>>> for Value {
             AnyValue::Null => *QNIL,
             AnyValue::Boolean(v) => Value::from(v),
             AnyValue::Utf8(v) => Value::from(v),
+            AnyValue::Date(v) => class::time()
+                .funcall::<_, _, Value>("at", (v * 86400,))
+                .unwrap()
+                .funcall::<_, _, Value>("utc", ())
+                .unwrap()
+                .funcall::<_, _, Value>("to_date", ())
+                .unwrap(),
             _ => todo!(),
         }
     }
@@ -102,10 +110,19 @@ impl TryConvert for Wrap<DataType> {
             "i32" => DataType::Int32,
             "i64" => DataType::Int64,
             "str" => DataType::Utf8,
+            "bin" => DataType::Binary,
             "bool" => DataType::Boolean,
+            "cat" => DataType::Categorical(None),
+            "date" => DataType::Date,
+            "datetime" => DataType::Datetime(TimeUnit::Microseconds, None),
             "f32" => DataType::Float32,
+            "time" => DataType::Time,
+            "dur" => DataType::Duration(TimeUnit::Microseconds),
             "f64" => DataType::Float64,
-            "date" => DataType::Date,
+            // "obj" => DataType::Object(OBJECT_NAME),
+            "list" => DataType::List(Box::new(DataType::Boolean)),
+            "null" => DataType::Null,
+            "unk" => DataType::Unknown,
             _ => {
                 return Err(RbValueError::new_err(format!(
                     "{} is not a supported DataType.",
@@ -407,3 +424,19 @@ pub fn parse_parquet_compression(
     };
     Ok(parsed)
 }
+
+pub struct ObjectValue {
+    pub inner: Value,
+}
+
+impl From<&dyn PolarsObjectSafe> for &ObjectValue {
+    fn from(val: &dyn PolarsObjectSafe) -> Self {
+        unsafe { &*(val as *const dyn PolarsObjectSafe as *const ObjectValue) }
+    }
+}
+
+impl ObjectValue {
+    pub fn to_object(&self) -> Value {
+        self.inner
+    }
+}

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/lazy/dataframe.rs

@@ -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)]
@@ -886,6 +887,19 @@ impl RbExpr {
         self.inner.clone().dt().round(&every, &offset).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()
     }
@@ -1305,6 +1319,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 +1361,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 +1398,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 +1425,14 @@ 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!();
+    }
+}
+
 #[magnus::wrap(class = "Polars::RbWhen")]
 #[derive(Clone)]
 pub struct RbWhen {