polars-df 0.1.0 → 0.1.1

data/ext/polars/src/series.rs

@@ -1,11 +1,12 @@
-use crate::conversion::wrap;
-use crate::{RbDataFrame, RbPolarsErr, RbResult};
 use magnus::exception::arg_error;
 use magnus::{Error, RArray, Value};
 use polars::prelude::*;
 use polars::series::IsSorted;
 use std::cell::RefCell;
 
+use crate::conversion::*;
+use crate::{RbDataFrame, RbPolarsErr, RbResult, RbValueError};
+
 #[magnus::wrap(class = "Polars::RbSeries")]
 pub struct RbSeries {
     pub series: RefCell<Series>,
@@ -24,6 +25,14 @@ impl RbSeries {
         }
     }
 
+    pub fn is_sorted_flag(&self) -> bool {
+        matches!(self.series.borrow().is_sorted(), IsSorted::Ascending)
+    }
+
+    pub fn is_sorted_reverse_flag(&self) -> bool {
+        matches!(self.series.borrow().is_sorted(), IsSorted::Descending)
+    }
+
     pub fn new_opt_bool(name: String, obj: RArray, strict: bool) -> RbResult<RbSeries> {
         let len = obj.len();
         let mut builder = BooleanChunkedBuilder::new(&name, len);
@@ -114,6 +123,29 @@ impl RbSeries {
         Ok(RbSeries::new(s))
     }
 
+    pub fn estimated_size(&self) -> usize {
+        self.series.borrow().estimated_size()
+    }
+
+    pub fn get_fmt(&self, index: usize, str_lengths: usize) -> String {
+        let val = format!("{}", self.series.borrow().get(index));
+        if let DataType::Utf8 | DataType::Categorical(_) = self.series.borrow().dtype() {
+            let v_trunc = &val[..val
+                .char_indices()
+                .take(str_lengths)
+                .last()
+                .map(|(i, c)| i + c.len_utf8())
+                .unwrap_or(0)];
+            if val == v_trunc {
+                val
+            } else {
+                format!("{}...", v_trunc)
+            }
+        } else {
+            val
+        }
+    }
+
     pub fn rechunk(&self, in_place: bool) -> Option<Self> {
         let series = self.series.borrow_mut().rechunk();
         if in_place {
@@ -124,6 +156,10 @@ impl RbSeries {
         }
     }
 
+    pub fn get_idx(&self, idx: usize) -> Value {
+        Wrap(self.series.borrow().get(idx)).into()
+    }
+
     pub fn bitand(&self, other: &RbSeries) -> RbResult<Self> {
         let out = self
             .series
@@ -196,15 +232,15 @@ impl RbSeries {
     }
 
     pub fn max(&self) -> Value {
-        wrap(self.series.borrow().max_as_series().get(0))
+        Wrap(self.series.borrow().max_as_series().get(0)).into()
     }
 
     pub fn min(&self) -> Value {
-        wrap(self.series.borrow().min_as_series().get(0))
+        Wrap(self.series.borrow().min_as_series().get(0)).into()
     }
 
     pub fn sum(&self) -> Value {
-        wrap(self.series.borrow().sum_as_series().get(0))
+        Wrap(self.series.borrow().sum_as_series().get(0)).into()
     }
 
     pub fn n_chunks(&self) -> usize {
@@ -454,6 +490,111 @@ impl RbSeries {
         }
     }
 
+    pub fn quantile(
+        &self,
+        quantile: f64,
+        interpolation: Wrap<QuantileInterpolOptions>,
+    ) -> RbResult<Value> {
+        Ok(Wrap(
+            self.series
+                .borrow()
+                .quantile_as_series(quantile, interpolation.0)
+                .map_err(|_| RbValueError::new_err("invalid quantile".into()))?
+                .get(0),
+        )
+        .into())
+    }
+
+    pub fn clone(&self) -> Self {
+        RbSeries::new(self.series.borrow().clone())
+    }
+
+    pub fn zip_with(&self, mask: &RbSeries, other: &RbSeries) -> RbResult<Self> {
+        let binding = mask.series.borrow();
+        let mask = binding.bool().map_err(RbPolarsErr::from)?;
+        let s = self
+            .series
+            .borrow()
+            .zip_with(mask, &other.series.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbSeries::new(s))
+    }
+
+    pub fn to_dummies(&self) -> RbResult<RbDataFrame> {
+        let df = self
+            .series
+            .borrow()
+            .to_dummies()
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn peak_max(&self) -> Self {
+        self.series.borrow().peak_max().into_series().into()
+    }
+
+    pub fn peak_min(&self) -> Self {
+        self.series.borrow().peak_min().into_series().into()
+    }
+
+    pub fn n_unique(&self) -> RbResult<usize> {
+        let n = self.series.borrow().n_unique().map_err(RbPolarsErr::from)?;
+        Ok(n)
+    }
+
+    pub fn floor(&self) -> RbResult<Self> {
+        let s = self.series.borrow().floor().map_err(RbPolarsErr::from)?;
+        Ok(s.into())
+    }
+
+    pub fn shrink_to_fit(&self) {
+        self.series.borrow_mut().shrink_to_fit();
+    }
+
+    pub fn dot(&self, other: &RbSeries) -> Option<f64> {
+        self.series.borrow().dot(&other.series.borrow())
+    }
+
+    pub fn skew(&self, bias: bool) -> RbResult<Option<f64>> {
+        let out = self.series.borrow().skew(bias).map_err(RbPolarsErr::from)?;
+        Ok(out)
+    }
+
+    pub fn kurtosis(&self, fisher: bool, bias: bool) -> RbResult<Option<f64>> {
+        let out = self
+            .series
+            .borrow()
+            .kurtosis(fisher, bias)
+            .map_err(RbPolarsErr::from)?;
+        Ok(out)
+    }
+
+    pub fn cast(&self, dtype: Wrap<DataType>, strict: bool) -> RbResult<Self> {
+        let dtype = dtype.0;
+        let out = if strict {
+            self.series.borrow().strict_cast(&dtype)
+        } else {
+            self.series.borrow().cast(&dtype)
+        };
+        let out = out.map_err(RbPolarsErr::from)?;
+        Ok(out.into())
+    }
+
+    pub fn time_unit(&self) -> Option<String> {
+        if let DataType::Datetime(tu, _) | DataType::Duration(tu) = self.series.borrow().dtype() {
+            Some(
+                match tu {
+                    TimeUnit::Nanoseconds => "ns",
+                    TimeUnit::Microseconds => "us",
+                    TimeUnit::Milliseconds => "ms",
+                }
+                .to_string(),
+            )
+        } else {
+            None
+        }
+    }
+
     // dispatch dynamically in future?
 
     pub fn cumsum(&self, reverse: bool) -> Self {
@@ -468,8 +609,30 @@ impl RbSeries {
         self.series.borrow().cummin(reverse).into()
     }
 
+    pub fn cumprod(&self, reverse: bool) -> Self {
+        self.series.borrow().cumprod(reverse).into()
+    }
+
     pub fn slice(&self, offset: i64, length: usize) -> Self {
         let series = self.series.borrow().slice(offset, length);
         series.into()
     }
+
+    pub fn ceil(&self) -> RbResult<Self> {
+        let s = self.series.borrow().ceil().map_err(RbPolarsErr::from)?;
+        Ok(s.into())
+    }
+
+    pub fn round(&self, decimals: u32) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .round(decimals)
+            .map_err(RbPolarsErr::from)?;
+        Ok(s.into())
+    }
+}
+
+pub fn to_rbseries_collection(s: Vec<Series>) -> Vec<RbSeries> {
+    s.into_iter().map(|v| RbSeries::new(v)).collect()
 }

data/lib/polars/cat_expr.rb

@@ -0,0 +1,13 @@
+module Polars
+  class CatExpr
+    attr_accessor :_rbexpr
+
+    def initialize(expr)
+      self._rbexpr = expr._rbexpr
+    end
+
+    def set_ordering(ordering)
+      Utils.wrap_expr(_rbexpr.cat_set_ordering(ordering))
+    end
+  end
+end

data/lib/polars/data_frame.rb

@@ -79,10 +79,51 @@ module Polars
       _df.columns
     end
 
+    def columns=(columns)
+      _df.set_column_names(columns)
+    end
+
     def dtypes
       _df.dtypes.map(&:to_sym)
     end
 
+    def schema
+      columns.zip(dtypes).to_h
+    end
+
+    # def ==(other)
+    # end
+
+    # def !=(other)
+    # end
+
+    # def >(other)
+    # end
+
+    # def <(other)
+    # end
+
+    # def >=(other)
+    # end
+
+    # def <=(other)
+    # end
+
+    # def *(other)
+    # end
+
+    # def /(other)
+    # end
+
+    # def +(other)
+    # end
+
+    # def -(other)
+    # end
+
+    # def %(other)
+    # end
+
     def to_s
       _df.to_s
     end
@@ -96,6 +137,25 @@ module Polars
       Utils.wrap_s(_df.column(name))
     end
 
+    # def []=(key, value)
+    # end
+
+    def to_h(as_series: true)
+      if as_series
+        get_columns.to_h { |s| [s.name, s] }
+      else
+        get_columns.to_h { |s| [s.name, s.to_a] }
+      end
+    end
+
+    # def to_hs / to_a
+    # end
+
+    # def to_numo
+    # end
+
+    # no to_pandas
+
     def to_series(index = 0)
       if index < 0
         index = columns.length + index
@@ -183,6 +243,12 @@ module Polars
       nil
     end
 
+    # def write_avro
+    # end
+
+    # def write_ipc
+    # end
+
     def write_parquet(
       file,
       compression: "zstd",
@@ -202,10 +268,43 @@ module Polars
       )
     end
 
+    def estimated_size(unit = "b")
+      sz = _df.estimated_size
+      Utils.scale_bytes(sz, to: unit)
+    end
+
+    # def transpose
+    # end
+
+    def reverse
+      select(Polars.col("*").reverse)
+    end
+
+    def rename(mapping)
+      lazy.rename(mapping).collect(no_optimization: true)
+    end
+
+    def insert_at_idx(index, series)
+      if index < 0
+        index = columns.length + index
+      end
+      _df.insert_at_idx(index, series._s)
+      self
+    end
+
     def filter(predicate)
       lazy.filter(predicate).collect
     end
 
+    # def describe
+    # end
+
+    # def find_idx_by_name
+    # end
+
+    # def replace_at_idx
+    # end
+
     def sort(by, reverse: false, nulls_last: false)
       _from_rbdf(_df.sort(by, reverse, nulls_last))
     end
@@ -214,6 +313,16 @@ module Polars
       _df.frame_equal(other._df, null_equal)
     end
 
+    # def replace
+    # end
+
+    def slice(offset, length = nil)
+      if !length.nil? && length < 0
+        length = height - offset + length
+      end
+      _from_rbdf(_df.slice(offset, length))
+    end
+
     def limit(n = 5)
       head(n)
     end
@@ -226,10 +335,31 @@ module Polars
       _from_rbdf(_df.tail(n))
     end
 
+    # def drop_nulls
+    # end
+
+    # def pipe
+    # end
+
+    # def with_row_count
+    # end
+
     def groupby(by, maintain_order: false)
       lazy.groupby(by, maintain_order: maintain_order)
     end
 
+    # def groupby_rolling
+    # end
+
+    # def groupby_dynamic
+    # end
+
+    # def upsample
+    # end
+
+    # def join_asof
+    # end
+
     def join(other, left_on: nil, right_on: nil, on: nil, how: "inner", suffix: "_right")
       lazy
         .join(
@@ -243,12 +373,79 @@ module Polars
         .collect(no_optimization: true)
     end
 
+    # def apply
+    # end
+
     def with_column(column)
       lazy
         .with_column(column)
         .collect(no_optimization: true, string_cache: false)
     end
 
+    # def hstack
+    # end
+
+    # def vstack
+    # end
+
+    # def extend
+    # end
+
+    # def drop
+    # end
+
+    # def drop_in_place
+    # end
+
+    # def cleared
+    # end
+
+    # clone handled by initialize_copy
+
+    def get_columns
+      _df.get_columns.map { |s| Utils.wrap_s(s) }
+    end
+
+    def get_column(name)
+      self[name]
+    end
+
+    # def fill_null
+    # end
+
+    def fill_nan(fill_value)
+      lazy.fill_nan(fill_value).collect(no_optimization: true)
+    end
+
+    # def explode
+    # end
+
+    # def pivot
+    # end
+
+    # def melt
+    # end
+
+    # def unstack
+    # end
+
+    # def partition_by
+    # end
+
+    # def shift
+    # end
+
+    # def shift_and_fill
+    # end
+
+    def is_duplicated
+      Utils.wrap_s(_df.is_duplicated)
+    end
+
+    def is_unique
+      Utils.wrap_s(_df.is_unique)
+    end
+
     def lazy
       wrap_ldf(_df.lazy)
     end
@@ -262,6 +459,56 @@ module Polars
       )
     end
 
+    def with_columns(exprs)
+      if !exprs.nil? && !exprs.is_a?(Array)
+        exprs = [exprs]
+      end
+      lazy
+        .with_columns(exprs)
+        .collect(no_optimization: true, string_cache: false)
+    end
+
+    def n_chunks(strategy: "first")
+      if strategy == "first"
+        _df.n_chunks
+      elsif strategy == "all"
+        get_columns.map(&:n_chunks)
+      else
+        raise ArgumentError, "Strategy: '{strategy}' not understood. Choose one of {{'first',  'all'}}"
+      end
+    end
+
+    def max(axis: 0)
+      if axis == 0
+        _from_rbdf(_df.max)
+      elsif axis == 1
+        Utils.wrap_s(_df.hmax)
+      else
+        raise ArgumentError, "Axis should be 0 or 1."
+      end
+    end
+
+    def min(axis: 0)
+      if axis == 0
+        _from_rbdf(_df.min)
+      elsif axis == 1
+        Utils.wrap_s(_df.hmin)
+      else
+        raise ArgumentError, "Axis should be 0 or 1."
+      end
+    end
+
+    def sum(axis: 0, null_strategy: "ignore")
+      case axis
+      when 0
+        _from_rbdf(_df.sum)
+      when 1
+        Utils.wrap_s(_df.hsum(null_strategy))
+      else
+        raise ArgumentError, "Axis should be 0 or 1."
+      end
+    end
+
     def mean(axis: 0, null_strategy: "ignore")
       case axis
       when 0
@@ -273,15 +520,33 @@ module Polars
       end
     end
 
-    def with_columns(exprs)
-      if !exprs.nil? && !exprs.is_a?(Array)
-        exprs = [exprs]
-      end
-      lazy
-        .with_columns(exprs)
-        .collect(no_optimization: true, string_cache: false)
+    def std(ddof: 1)
+      _from_rbdf(_df.std(ddof))
+    end
+
+    def var(ddof: 1)
+      _from_rbdf(_df.var(ddof))
     end
 
+    def median
+      _from_rbdf(_df.median)
+    end
+
+    # def product
+    # end
+
+    # def quantile(quantile, interpolation: "nearest")
+    # end
+
+    # def to_dummies
+    # end
+
+    # def unique
+    # end
+
+    # def n_unique
+    # end
+
     def rechunk
       _from_rbdf(_df.rechunk)
     end
@@ -290,8 +555,48 @@ module Polars
       _from_rbdf(_df.null_count)
     end
 
+    # def sample
+    # end
+
+    # def fold
+    # end
+
+    # def row
+    # end
+
+    # def rows
+    # end
+
+    # def shrink_to_fit
+    # end
+
+    # def take_every
+    # end
+
+    # def hash_rows
+    # end
+
+    # def interpolate
+    # end
+
+    def is_empty
+      height == 0
+    end
+    alias_method :empty?, :is_empty
+
+    # def to_struct(name)
+    # end
+
+    # def unnest
+    # end
+
     private
 
+    def initialize_copy(other)
+      super
+      self._df = _df._clone
+    end
+
     def hash_to_rbdf(data)
       RbDataFrame.read_hash(data)
     end