polars-df 0.1.0 → 0.1.2

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);
@@ -107,11 +116,33 @@ init_method_opt!(new_opt_f32, Float32Type, f32);
 init_method_opt!(new_opt_f64, Float64Type, f64);
 
 impl RbSeries {
-    pub fn new_str(name: String, val: RArray, _strict: bool) -> RbResult<Self> {
-        let v = val.try_convert::<Vec<Option<String>>>()?;
-        let mut s = Utf8Chunked::new(&name, v).into_series();
+    pub fn new_str(name: String, val: Wrap<Utf8Chunked>, _strict: bool) -> Self {
+        let mut s = val.0.into_series();
         s.rename(&name);
-        Ok(RbSeries::new(s))
+        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> {
@@ -124,6 +155,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
@@ -163,16 +198,16 @@ impl RbSeries {
         self.series.borrow_mut().rename(&name);
     }
 
-    pub fn dtype(&self) -> String {
-        self.series.borrow().dtype().to_string()
+    pub fn dtype(&self) -> Value {
+        Wrap(self.series.borrow().dtype().clone()).into()
     }
 
-    pub fn inner_dtype(&self) -> Option<String> {
+    pub fn inner_dtype(&self) -> Option<Value> {
         self.series
             .borrow()
             .dtype()
             .inner_dtype()
-            .map(|dt| dt.to_string())
+            .map(|dt| Wrap(dt.clone()).into())
     }
 
     pub fn set_sorted(&self, reverse: bool) -> Self {
@@ -196,15 +231,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 +489,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 +608,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(RbSeries::new).collect()
 }

data/lib/polars/batched_csv_reader.rb

@@ -0,0 +1,95 @@
+module Polars
+  class BatchedCsvReader
+    attr_accessor :_reader, :new_columns
+
+    def initialize(
+      file,
+      has_header: true,
+      columns: nil,
+      sep: ",",
+      comment_char: nil,
+      quote_char: '"',
+      skip_rows: 0,
+      dtypes: nil,
+      null_values: nil,
+      ignore_errors: false,
+      parse_dates: false,
+      n_threads: nil,
+      infer_schema_length: 100,
+      batch_size: 50_000,
+      n_rows: nil,
+      encoding: "utf8",
+      low_memory: false,
+      rechunk: true,
+      skip_rows_after_header: 0,
+      row_count_name: nil,
+      row_count_offset: 0,
+      sample_size: 1024,
+      eol_char: "\n",
+      new_columns: nil
+    )
+      if file.is_a?(String) || (defined?(Pathname) && file.is_a?(Pathname))
+        path = Utils.format_path(file)
+      end
+
+      dtype_list = nil
+      dtype_slice = nil
+      if !dtypes.nil?
+        if dtypes.is_a?(Hash)
+          dtype_list = []
+          dtypes.each do|k, v|
+            dtype_list << [k, Utils.rb_type_to_dtype(v)]
+          end
+        elsif dtypes.is_a?(Array)
+          dtype_slice = dtypes
+        else
+          raise ArgumentError, "dtype arg should be list or dict"
+        end
+      end
+
+      processed_null_values = Utils._process_null_values(null_values)
+      projection, columns = Utils.handle_projection_columns(columns)
+
+      self._reader = RbBatchedCsv.new(
+        infer_schema_length,
+        batch_size,
+        has_header,
+        ignore_errors,
+        n_rows,
+        skip_rows,
+        projection,
+        sep,
+        rechunk,
+        columns,
+        encoding,
+        n_threads,
+        path,
+        dtype_list,
+        dtype_slice,
+        low_memory,
+        comment_char,
+        quote_char,
+        processed_null_values,
+        parse_dates,
+        skip_rows_after_header,
+        Utils._prepare_row_count_args(row_count_name, row_count_offset),
+        sample_size,
+        eol_char
+      )
+      self.new_columns = new_columns
+    end
+
+    def next_batches(n)
+      batches = _reader.next_batches(n)
+      if !batches.nil?
+        if new_columns
+          batches.map { |df| Utils._update_columns(Utils.wrap_df(df), new_columns) }
+        else
+          batches.map { |df| Utils.wrap_df(df) }
+        end
+      else
+        nil
+      end
+    end
+  end
+end

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