polars-df 0.1.0 → 0.1.1

data/Cargo.toml

@@ -0,0 +1,5 @@
+[workspace]
+members = ["ext/polars"]
+
+[profile.release]
+strip = true

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.1.0"
+version = "0.1.1"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
 publish = false
@@ -15,21 +15,51 @@ serde_json = "1"
 [dependencies.polars]
 version = "0.25.1"
 features = [
+    "abs",
     "arange",
+    "concat_str",
     "csv-file",
     "cum_agg",
+    "cumulative_eval",
+    "date_offset",
     "diagonal_concat",
+    "diff",
+    "dot_product",
     "dtype-full",
+    "dynamic_groupby",
+    "ewma",
     "fmt",
     "horizontal_concat",
     "interpolate",
     "ipc",
+    "is_first",
+    "is_in",
     "json",
     "lazy",
     "lazy_regex",
+    "list_eval",
+    "log",
+    "meta",
+    "mode",
+    "moment",
     "parquet",
+    "partition_by",
+    "pct_change",
+    "product",
+    "random",
+    "rank",
+    "repeat_by",
+    "rolling_window",
+    "round_series",
+    "search_sorted",
     "semi_anti_join",
     "serde-lazy",
+    "sign",
+    "string_justify",
     "strings",
+    "timezones",
+    "to_dummies",
+    "top_k",
     "trigonometry",
+    "unique_counts",
 ]

data/ext/polars/src/conversion.rs

@@ -1,27 +1,245 @@
-use magnus::{Value, QNIL};
+use magnus::{TryConvert, Value, QNIL};
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
 use polars::datatypes::AnyValue;
 use polars::frame::DataFrame;
 use polars::prelude::*;
+use polars::series::ops::NullBehavior;
 
-use crate::{RbDataFrame, RbResult, RbValueError};
-
-pub fn wrap(val: AnyValue) -> Value {
-    match val {
-        AnyValue::UInt8(v) => Value::from(v),
-        AnyValue::UInt16(v) => Value::from(v),
-        AnyValue::UInt32(v) => Value::from(v),
-        AnyValue::UInt64(v) => Value::from(v),
-        AnyValue::Int8(v) => Value::from(v),
-        AnyValue::Int16(v) => Value::from(v),
-        AnyValue::Int32(v) => Value::from(v),
-        AnyValue::Int64(v) => Value::from(v),
-        AnyValue::Float32(v) => Value::from(v),
-        AnyValue::Float64(v) => Value::from(v),
-        AnyValue::Null => *QNIL,
-        AnyValue::Boolean(v) => Value::from(v),
-        AnyValue::Utf8(v) => Value::from(v),
-        _ => todo!(),
+use crate::{RbDataFrame, RbPolarsErr, RbResult, RbValueError};
+
+pub struct Wrap<T>(pub T);
+
+impl<T> From<T> for Wrap<T> {
+    fn from(t: T) -> Self {
+        Wrap(t)
+    }
+}
+
+pub fn get_df(obj: Value) -> RbResult<DataFrame> {
+    let rbdf = obj.funcall::<_, _, &RbDataFrame>("_df", ())?;
+    Ok(rbdf.df.borrow().clone())
+}
+
+impl Into<Value> for Wrap<AnyValue<'_>> {
+    fn into(self) -> Value {
+        match self.0 {
+            AnyValue::UInt8(v) => Value::from(v),
+            AnyValue::UInt16(v) => Value::from(v),
+            AnyValue::UInt32(v) => Value::from(v),
+            AnyValue::UInt64(v) => Value::from(v),
+            AnyValue::Int8(v) => Value::from(v),
+            AnyValue::Int16(v) => Value::from(v),
+            AnyValue::Int32(v) => Value::from(v),
+            AnyValue::Int64(v) => Value::from(v),
+            AnyValue::Float32(v) => Value::from(v),
+            AnyValue::Float64(v) => Value::from(v),
+            AnyValue::Null => *QNIL,
+            AnyValue::Boolean(v) => Value::from(v),
+            AnyValue::Utf8(v) => Value::from(v),
+            _ => todo!(),
+        }
+    }
+}
+
+impl TryConvert for Wrap<DataType> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let dtype = match ob.try_convert::<String>()?.as_str() {
+            "u8" => DataType::UInt8,
+            "u16" => DataType::UInt16,
+            "u32" => DataType::UInt32,
+            "u64" => DataType::UInt64,
+            "i8" => DataType::Int8,
+            "i16" => DataType::Int16,
+            "i32" => DataType::Int32,
+            "i64" => DataType::Int64,
+            "str" => DataType::Utf8,
+            "bool" => DataType::Boolean,
+            "f32" => DataType::Float32,
+            "f64" => DataType::Float64,
+            "date" => DataType::Date,
+            _ => {
+                return Err(RbValueError::new_err(format!(
+                    "{} is not a supported DataType.",
+                    ob
+                )))
+            }
+        };
+        Ok(Wrap(dtype))
+    }
+}
+
+impl<'s> TryConvert for Wrap<AnyValue<'s>> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        // TODO improve
+        if let Ok(v) = ob.try_convert::<i64>() {
+            Ok(AnyValue::Int64(v).into())
+        } else if let Ok(v) = ob.try_convert::<f64>() {
+            Ok(AnyValue::Float64(v).into())
+        } else {
+            Err(RbPolarsErr::other(format!(
+                "object type not supported {:?}",
+                ob
+            )))
+        }
+    }
+}
+
+impl TryConvert for Wrap<CategoricalOrdering> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "physical" => CategoricalOrdering::Physical,
+            "lexical" => CategoricalOrdering::Lexical,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "ordering must be one of {{'physical', 'lexical'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<ClosedWindow> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "left" => ClosedWindow::Left,
+            "right" => ClosedWindow::Right,
+            "both" => ClosedWindow::Both,
+            "none" => ClosedWindow::None,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "closed must be one of {{'left', 'right', 'both', 'none'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<JoinType> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "inner" => JoinType::Inner,
+            "left" => JoinType::Left,
+            "outer" => JoinType::Outer,
+            "semi" => JoinType::Semi,
+            "anti" => JoinType::Anti,
+            // #[cfg(feature = "cross_join")]
+            // "cross" => JoinType::Cross,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                "how must be one of {{'inner', 'left', 'outer', 'semi', 'anti', 'cross'}}, got {}",
+                v
+            )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<NullBehavior> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "drop" => NullBehavior::Drop,
+            "ignore" => NullBehavior::Ignore,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "null behavior must be one of {{'drop', 'ignore'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<NullStrategy> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "ignore" => NullStrategy::Ignore,
+            "propagate" => NullStrategy::Propagate,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "null strategy must be one of {{'ignore', 'propagate'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<QuantileInterpolOptions> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "lower" => QuantileInterpolOptions::Lower,
+            "higher" => QuantileInterpolOptions::Higher,
+            "nearest" => QuantileInterpolOptions::Nearest,
+            "linear" => QuantileInterpolOptions::Linear,
+            "midpoint" => QuantileInterpolOptions::Midpoint,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "interpolation must be one of {{'lower', 'higher', 'nearest', 'linear', 'midpoint'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<RankMethod> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "min" => RankMethod::Min,
+            "max" => RankMethod::Max,
+            "average" => RankMethod::Average,
+            "dense" => RankMethod::Dense,
+            "ordinal" => RankMethod::Ordinal,
+            "random" => RankMethod::Random,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "method must be one of {{'min', 'max', 'average', 'dense', 'ordinal', 'random'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<TimeUnit> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "ns" => TimeUnit::Nanoseconds,
+            "us" => TimeUnit::Microseconds,
+            "ms" => TimeUnit::Milliseconds,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "time unit must be one of {{'ns', 'us', 'ms'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<UniqueKeepStrategy> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "first" => UniqueKeepStrategy::First,
+            "last" => UniqueKeepStrategy::Last,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "keep must be one of {{'first', 'last'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
     }
 }
 
@@ -47,30 +265,6 @@ pub fn parse_fill_null_strategy(
     Ok(parsed)
 }
 
-pub fn wrap_join_type(ob: &str) -> RbResult<JoinType> {
-    let parsed = match ob {
-        "inner" => JoinType::Inner,
-        "left" => JoinType::Left,
-        "outer" => JoinType::Outer,
-        "semi" => JoinType::Semi,
-        "anti" => JoinType::Anti,
-        // #[cfg(feature = "cross_join")]
-        // "cross" => JoinType::Cross,
-        v => {
-            return Err(RbValueError::new_err(format!(
-                "how must be one of {{'inner', 'left', 'outer', 'semi', 'anti', 'cross'}}, got {}",
-                v
-            )))
-        }
-    };
-    Ok(parsed)
-}
-
-pub fn get_df(obj: Value) -> RbResult<DataFrame> {
-    let rbdf = obj.funcall::<_, _, &RbDataFrame>("_df", ())?;
-    Ok(rbdf.df.borrow().clone())
-}
-
 pub fn parse_parquet_compression(
     compression: &str,
     compression_level: Option<i32>,

data/ext/polars/src/dataframe.rs

@@ -7,8 +7,9 @@ use std::io::{BufReader, BufWriter, Cursor};
 use std::ops::Deref;
 use std::path::PathBuf;
 
-use crate::conversion::parse_parquet_compression;
+use crate::conversion::*;
 use crate::file::{get_file_like, get_mmap_bytes_reader};
+use crate::series::to_rbseries_collection;
 use crate::{series, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
 
 #[magnus::wrap(class = "Polars::RbDataFrame")]
@@ -38,6 +39,10 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
+    pub fn estimated_size(&self) -> usize {
+        self.df.borrow().estimated_size()
+    }
+
     pub fn read_csv(rb_f: Value, has_header: bool) -> RbResult<Self> {
         let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
         let df = CsvReader::new(mmap_bytes_r)
@@ -213,6 +218,11 @@ impl RbDataFrame {
         format!("{}", self.df.borrow())
     }
 
+    pub fn get_columns(&self) -> Vec<RbSeries> {
+        let cols = self.df.borrow().get_columns().clone();
+        to_rbseries_collection(cols)
+    }
+
     pub fn columns(&self) -> Vec<String> {
         self.df
             .borrow()
@@ -222,6 +232,14 @@ impl RbDataFrame {
             .collect()
     }
 
+    pub fn set_column_names(&self, names: Vec<String>) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .set_column_names(&names)
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
     pub fn dtypes(&self) -> Vec<String> {
         self.df
             .borrow()
@@ -230,6 +248,11 @@ impl RbDataFrame {
             .collect()
     }
 
+    pub fn n_chunks(&self) -> RbResult<usize> {
+        let n = self.df.borrow().n_chunks().map_err(RbPolarsErr::from)?;
+        Ok(n)
+    }
+
     pub fn shape(&self) -> (usize, usize) {
         self.df.borrow().shape()
     }
@@ -258,6 +281,28 @@ impl RbDataFrame {
             .map_err(RbPolarsErr::from)
     }
 
+    pub fn select(&self, selection: Vec<String>) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .select(selection)
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
+    pub fn take(&self, indices: Vec<IdxSize>) -> RbResult<Self> {
+        let indices = IdxCa::from_vec("", indices);
+        let df = self.df.borrow().take(&indices).map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
+    pub fn take_with_series(&self, indices: &RbSeries) -> RbResult<Self> {
+        let binding = indices.series.borrow();
+        let idx = binding.idx().map_err(RbPolarsErr::from)?;
+        let df = self.df.borrow().take(idx).map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
     pub fn sort(&self, by_column: String, reverse: bool, nulls_last: bool) -> RbResult<Self> {
         let df = self
             .df
@@ -273,6 +318,38 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
+    pub fn replace(&self, column: String, new_col: &RbSeries) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .replace(&column, new_col.series.borrow().clone())
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn replace_at_idx(&self, index: usize, new_col: &RbSeries) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .replace_at_idx(index, new_col.series.borrow().clone())
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn insert_at_idx(&self, index: usize, new_col: &RbSeries) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .insert_at_idx(index, new_col.series.borrow().clone())
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn slice(&self, offset: usize, length: Option<usize>) -> Self {
+        let df = self.df.borrow().slice(
+            offset as i64,
+            length.unwrap_or_else(|| self.df.borrow().height()),
+        );
+        df.into()
+    }
+
     pub fn head(&self, length: Option<usize>) -> Self {
         self.df.borrow().head(length).into()
     }
@@ -281,6 +358,20 @@ impl RbDataFrame {
         self.df.borrow().tail(length).into()
     }
 
+    pub fn is_unique(&self) -> RbResult<RbSeries> {
+        let mask = self.df.borrow().is_unique().map_err(RbPolarsErr::from)?;
+        Ok(mask.into_series().into())
+    }
+
+    pub fn is_duplicated(&self) -> RbResult<RbSeries> {
+        let mask = self
+            .df
+            .borrow()
+            .is_duplicated()
+            .map_err(RbPolarsErr::from)?;
+        Ok(mask.into_series().into())
+    }
+
     pub fn frame_equal(&self, other: &RbDataFrame, null_equal: bool) -> bool {
         if null_equal {
             self.df.borrow().frame_equal_missing(&other.df.borrow())
@@ -289,16 +380,202 @@ impl RbDataFrame {
         }
     }
 
+    pub fn with_row_count(&self, name: String, offset: Option<IdxSize>) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .with_row_count(&name, offset)
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn clone(&self) -> Self {
+        RbDataFrame::new(self.df.borrow().clone())
+    }
+
+    pub fn melt(
+        &self,
+        id_vars: Vec<String>,
+        value_vars: Vec<String>,
+        value_name: Option<String>,
+        variable_name: Option<String>,
+    ) -> RbResult<Self> {
+        let args = MeltArgs {
+            id_vars,
+            value_vars,
+            value_name,
+            variable_name,
+        };
+
+        let df = self.df.borrow().melt2(args).map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
+    pub fn partition_by(&self, groups: Vec<String>, stable: bool) -> RbResult<Vec<Self>> {
+        let out = if stable {
+            self.df.borrow().partition_by_stable(groups)
+        } else {
+            self.df.borrow().partition_by(groups)
+        }
+        .map_err(RbPolarsErr::from)?;
+        Ok(out.into_iter().map(|v| RbDataFrame::new(v)).collect())
+    }
+
+    pub fn shift(&self, periods: i64) -> Self {
+        self.df.borrow().shift(periods).into()
+    }
+
+    pub fn unique(
+        &self,
+        maintain_order: bool,
+        subset: Option<Vec<String>>,
+        keep: Wrap<UniqueKeepStrategy>,
+    ) -> RbResult<Self> {
+        let subset = subset.as_ref().map(|v| v.as_ref());
+        let df = match maintain_order {
+            true => self.df.borrow().unique_stable(subset, keep.0),
+            false => self.df.borrow().unique(subset, keep.0),
+        }
+        .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
     pub fn lazy(&self) -> RbLazyFrame {
         self.df.borrow().clone().lazy().into()
     }
 
+    pub fn max(&self) -> Self {
+        self.df.borrow().max().into()
+    }
+
+    pub fn min(&self) -> Self {
+        self.df.borrow().min().into()
+    }
+
+    pub fn sum(&self) -> Self {
+        self.df.borrow().sum().into()
+    }
+
     pub fn mean(&self) -> Self {
         self.df.borrow().mean().into()
     }
 
+    pub fn std(&self, ddof: u8) -> Self {
+        self.df.borrow().std(ddof).into()
+    }
+
+    pub fn var(&self, ddof: u8) -> Self {
+        self.df.borrow().var(ddof).into()
+    }
+
+    pub fn median(&self) -> Self {
+        self.df.borrow().median().into()
+    }
+
+    pub fn hmean(&self, null_strategy: Wrap<NullStrategy>) -> RbResult<Option<RbSeries>> {
+        let s = self
+            .df
+            .borrow()
+            .hmean(null_strategy.0)
+            .map_err(RbPolarsErr::from)?;
+        Ok(s.map(|s| s.into()))
+    }
+
+    pub fn hmax(&self) -> RbResult<Option<RbSeries>> {
+        let s = self.df.borrow().hmax().map_err(RbPolarsErr::from)?;
+        Ok(s.map(|s| s.into()))
+    }
+
+    pub fn hmin(&self) -> RbResult<Option<RbSeries>> {
+        let s = self.df.borrow().hmin().map_err(RbPolarsErr::from)?;
+        Ok(s.map(|s| s.into()))
+    }
+
+    pub fn hsum(&self, null_strategy: Wrap<NullStrategy>) -> RbResult<Option<RbSeries>> {
+        let s = self
+            .df
+            .borrow()
+            .hsum(null_strategy.0)
+            .map_err(RbPolarsErr::from)?;
+        Ok(s.map(|s| s.into()))
+    }
+
+    pub fn quantile(
+        &self,
+        quantile: f64,
+        interpolation: Wrap<QuantileInterpolOptions>,
+    ) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .quantile(quantile, interpolation.0)
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn to_dummies(&self, columns: Option<Vec<String>>) -> RbResult<Self> {
+        let df = match columns {
+            Some(cols) => self
+                .df
+                .borrow()
+                .columns_to_dummies(cols.iter().map(|x| x as &str).collect()),
+            None => self.df.borrow().to_dummies(),
+        }
+        .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
     pub fn null_count(&self) -> Self {
         let df = self.df.borrow().null_count();
         df.into()
     }
+
+    pub fn shrink_to_fit(&self) {
+        self.df.borrow_mut().shrink_to_fit();
+    }
+
+    pub fn transpose(&self, include_header: bool, names: String) -> RbResult<Self> {
+        let mut df = self.df.borrow().transpose().map_err(RbPolarsErr::from)?;
+        if include_header {
+            let s = Utf8Chunked::from_iter_values(
+                &names,
+                self.df.borrow().get_columns().iter().map(|s| s.name()),
+            )
+            .into_series();
+            df.insert_at_idx(0, s).unwrap();
+        }
+        Ok(df.into())
+    }
+
+    pub fn upsample(
+        &self,
+        by: Vec<String>,
+        index_column: String,
+        every: String,
+        offset: String,
+        stable: bool,
+    ) -> RbResult<Self> {
+        let out = if stable {
+            self.df.borrow().upsample_stable(
+                by,
+                &index_column,
+                Duration::parse(&every),
+                Duration::parse(&offset),
+            )
+        } else {
+            self.df.borrow().upsample(
+                by,
+                &index_column,
+                Duration::parse(&every),
+                Duration::parse(&offset),
+            )
+        };
+        let out = out.map_err(RbPolarsErr::from)?;
+        Ok(out.into())
+    }
+
+    pub fn unnest(&self, names: Vec<String>) -> RbResult<Self> {
+        let df = self.df.borrow().unnest(names).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
 }