polars-df 0.1.1 → 0.1.3

data/ext/polars/src/series.rs

@@ -5,6 +5,7 @@ use polars::series::IsSorted;
 use std::cell::RefCell;
 
 use crate::conversion::*;
+use crate::set::set_at_idx;
 use crate::{RbDataFrame, RbPolarsErr, RbResult, RbValueError};
 
 #[magnus::wrap(class = "Polars::RbSeries")]
@@ -116,11 +117,10 @@ 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 {
@@ -199,16 +199,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 {
@@ -475,6 +475,12 @@ impl RbSeries {
             s.into_iter().collect()
         } else if let Ok(s) = series.utf8() {
             s.into_iter().collect()
+        } else if let Ok(_s) = series.date() {
+            let a = RArray::with_capacity(series.len());
+            for v in series.iter() {
+                a.push::<Value>(Wrap(v).into()).unwrap();
+            }
+            a
         } else {
             unimplemented!();
         }
@@ -595,44 +601,190 @@ impl RbSeries {
         }
     }
 
-    // dispatch dynamically in future?
-
-    pub fn cumsum(&self, reverse: bool) -> Self {
-        self.series.borrow().cumsum(reverse).into()
+    pub fn set_at_idx(&self, idx: &RbSeries, values: &RbSeries) -> RbResult<()> {
+        let mut s = self.series.borrow_mut();
+        match set_at_idx(s.clone(), &idx.series.borrow(), &values.series.borrow()) {
+            Ok(out) => {
+                *s = out;
+                Ok(())
+            }
+            Err(e) => Err(RbPolarsErr::from(e)),
+        }
     }
+}
 
-    pub fn cummax(&self, reverse: bool) -> Self {
-        self.series.borrow().cummax(reverse).into()
-    }
+macro_rules! impl_eq_num {
+    ($name:ident, $type:ty) => {
+        impl RbSeries {
+            pub fn $name(&self, rhs: $type) -> RbResult<Self> {
+                let s = self.series.borrow().equal(rhs).map_err(RbPolarsErr::from)?;
+                Ok(RbSeries::new(s.into_series()))
+            }
+        }
+    };
+}
 
-    pub fn cummin(&self, reverse: bool) -> Self {
-        self.series.borrow().cummin(reverse).into()
-    }
+impl_eq_num!(eq_u8, u8);
+impl_eq_num!(eq_u16, u16);
+impl_eq_num!(eq_u32, u32);
+impl_eq_num!(eq_u64, u64);
+impl_eq_num!(eq_i8, i8);
+impl_eq_num!(eq_i16, i16);
+impl_eq_num!(eq_i32, i32);
+impl_eq_num!(eq_i64, i64);
+impl_eq_num!(eq_f32, f32);
+impl_eq_num!(eq_f64, f64);
+// impl_eq_num!(eq_str, &str);
+
+macro_rules! impl_neq_num {
+    ($name:ident, $type:ty) => {
+        impl RbSeries {
+            pub fn $name(&self, rhs: $type) -> RbResult<Self> {
+                let s = self
+                    .series
+                    .borrow()
+                    .not_equal(rhs)
+                    .map_err(RbPolarsErr::from)?;
+                Ok(RbSeries::new(s.into_series()))
+            }
+        }
+    };
+}
 
-    pub fn cumprod(&self, reverse: bool) -> Self {
-        self.series.borrow().cumprod(reverse).into()
-    }
+impl_neq_num!(neq_u8, u8);
+impl_neq_num!(neq_u16, u16);
+impl_neq_num!(neq_u32, u32);
+impl_neq_num!(neq_u64, u64);
+impl_neq_num!(neq_i8, i8);
+impl_neq_num!(neq_i16, i16);
+impl_neq_num!(neq_i32, i32);
+impl_neq_num!(neq_i64, i64);
+impl_neq_num!(neq_f32, f32);
+impl_neq_num!(neq_f64, f64);
+// impl_neq_num!(neq_str, &str);
+
+macro_rules! impl_gt_num {
+    ($name:ident, $type:ty) => {
+        impl RbSeries {
+            pub fn $name(&self, rhs: $type) -> RbResult<Self> {
+                let s = self.series.borrow().gt(rhs).map_err(RbPolarsErr::from)?;
+                Ok(RbSeries::new(s.into_series()))
+            }
+        }
+    };
+}
 
-    pub fn slice(&self, offset: i64, length: usize) -> Self {
-        let series = self.series.borrow().slice(offset, length);
-        series.into()
-    }
+impl_gt_num!(gt_u8, u8);
+impl_gt_num!(gt_u16, u16);
+impl_gt_num!(gt_u32, u32);
+impl_gt_num!(gt_u64, u64);
+impl_gt_num!(gt_i8, i8);
+impl_gt_num!(gt_i16, i16);
+impl_gt_num!(gt_i32, i32);
+impl_gt_num!(gt_i64, i64);
+impl_gt_num!(gt_f32, f32);
+impl_gt_num!(gt_f64, f64);
+// impl_gt_num!(gt_str, &str);
+
+macro_rules! impl_gt_eq_num {
+    ($name:ident, $type:ty) => {
+        impl RbSeries {
+            pub fn $name(&self, rhs: $type) -> RbResult<Self> {
+                let s = self.series.borrow().gt_eq(rhs).map_err(RbPolarsErr::from)?;
+                Ok(RbSeries::new(s.into_series()))
+            }
+        }
+    };
+}
 
-    pub fn ceil(&self) -> RbResult<Self> {
-        let s = self.series.borrow().ceil().map_err(RbPolarsErr::from)?;
-        Ok(s.into())
-    }
+impl_gt_eq_num!(gt_eq_u8, u8);
+impl_gt_eq_num!(gt_eq_u16, u16);
+impl_gt_eq_num!(gt_eq_u32, u32);
+impl_gt_eq_num!(gt_eq_u64, u64);
+impl_gt_eq_num!(gt_eq_i8, i8);
+impl_gt_eq_num!(gt_eq_i16, i16);
+impl_gt_eq_num!(gt_eq_i32, i32);
+impl_gt_eq_num!(gt_eq_i64, i64);
+impl_gt_eq_num!(gt_eq_f32, f32);
+impl_gt_eq_num!(gt_eq_f64, f64);
+// impl_gt_eq_num!(gt_eq_str, &str);
+
+macro_rules! impl_lt_num {
+    ($name:ident, $type:ty) => {
+        impl RbSeries {
+            pub fn $name(&self, rhs: $type) -> RbResult<RbSeries> {
+                let s = self.series.borrow().lt(rhs).map_err(RbPolarsErr::from)?;
+                Ok(RbSeries::new(s.into_series()))
+            }
+        }
+    };
+}
 
-    pub fn round(&self, decimals: u32) -> RbResult<Self> {
-        let s = self
-            .series
-            .borrow()
-            .round(decimals)
-            .map_err(RbPolarsErr::from)?;
-        Ok(s.into())
-    }
+impl_lt_num!(lt_u8, u8);
+impl_lt_num!(lt_u16, u16);
+impl_lt_num!(lt_u32, u32);
+impl_lt_num!(lt_u64, u64);
+impl_lt_num!(lt_i8, i8);
+impl_lt_num!(lt_i16, i16);
+impl_lt_num!(lt_i32, i32);
+impl_lt_num!(lt_i64, i64);
+impl_lt_num!(lt_f32, f32);
+impl_lt_num!(lt_f64, f64);
+// impl_lt_num!(lt_str, &str);
+
+macro_rules! impl_lt_eq_num {
+    ($name:ident, $type:ty) => {
+        impl RbSeries {
+            pub fn $name(&self, rhs: $type) -> RbResult<Self> {
+                let s = self.series.borrow().lt_eq(rhs).map_err(RbPolarsErr::from)?;
+                Ok(RbSeries::new(s.into_series()))
+            }
+        }
+    };
+}
+
+impl_lt_eq_num!(lt_eq_u8, u8);
+impl_lt_eq_num!(lt_eq_u16, u16);
+impl_lt_eq_num!(lt_eq_u32, u32);
+impl_lt_eq_num!(lt_eq_u64, u64);
+impl_lt_eq_num!(lt_eq_i8, i8);
+impl_lt_eq_num!(lt_eq_i16, i16);
+impl_lt_eq_num!(lt_eq_i32, i32);
+impl_lt_eq_num!(lt_eq_i64, i64);
+impl_lt_eq_num!(lt_eq_f32, f32);
+impl_lt_eq_num!(lt_eq_f64, f64);
+// impl_lt_eq_num!(lt_eq_str, &str);
+
+pub fn to_series_collection(rs: RArray) -> RbResult<Vec<Series>> {
+    let mut series = Vec::new();
+    for item in rs.each() {
+        series.push(item?.try_convert::<&RbSeries>()?.series.borrow().clone());
+    }
+    Ok(series)
 }
 
 pub fn to_rbseries_collection(s: Vec<Series>) -> Vec<RbSeries> {
-    s.into_iter().map(|v| RbSeries::new(v)).collect()
+    s.into_iter().map(RbSeries::new).collect()
+}
+
+impl RbSeries {
+    pub fn new_opt_date(name: String, values: RArray, _strict: Option<bool>) -> RbResult<Self> {
+        let len = values.len();
+        let mut builder = PrimitiveChunkedBuilder::<Int32Type>::new(&name, len);
+        for item in values.each() {
+            let v = item?;
+            if v.is_nil() {
+                builder.append_null();
+            } else {
+                // convert to DateTime for UTC
+                let v: Value = v.funcall("to_datetime", ())?;
+                let v: Value = v.funcall("to_time", ())?;
+                let v: Value = v.funcall("to_i", ())?;
+                // TODO use strict
+                builder.append_value(v.try_convert::<i32>()? / 86400);
+            }
+        }
+        let ca: ChunkedArray<Int32Type> = builder.finish();
+        Ok(ca.into_date().into_series().into())
+    }
 }

data/ext/polars/src/set.rs

@@ -0,0 +1,91 @@
+// use polars::export::arrow2::array::Array;
+use polars::prelude::*;
+
+pub fn set_at_idx(mut s: Series, idx: &Series, values: &Series) -> PolarsResult<Series> {
+    let logical_dtype = s.dtype().clone();
+    let idx = idx.cast(&IDX_DTYPE)?;
+    let idx = idx.rechunk();
+    let idx = idx.idx().unwrap();
+    let idx = idx.downcast_iter().next().unwrap();
+
+    // if idx.null_count() > 0 {
+    //     return Err(PolarsError::ComputeError(
+    //         "index values should not be null".into(),
+    //     ));
+    // }
+
+    let idx = idx.values().as_slice();
+
+    let values = values.to_physical_repr().cast(&s.dtype().to_physical())?;
+
+    // do not shadow, otherwise s is not dropped immediately
+    // and we want to have mutable access
+    s = s.to_physical_repr().into_owned();
+    let mutable_s = s._get_inner_mut();
+
+    let s = match logical_dtype.to_physical() {
+        DataType::Int8 => {
+            let ca: &mut ChunkedArray<Int8Type> = mutable_s.as_mut();
+            let values = values.i8()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::Int16 => {
+            let ca: &mut ChunkedArray<Int16Type> = mutable_s.as_mut();
+            let values = values.i16()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::Int32 => {
+            let ca: &mut ChunkedArray<Int32Type> = mutable_s.as_mut();
+            let values = values.i32()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::Int64 => {
+            let ca: &mut ChunkedArray<Int64Type> = mutable_s.as_mut();
+            let values = values.i64()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::UInt8 => {
+            let ca: &mut ChunkedArray<UInt8Type> = mutable_s.as_mut();
+            let values = values.u8()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::UInt16 => {
+            let ca: &mut ChunkedArray<UInt16Type> = mutable_s.as_mut();
+            let values = values.u16()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::UInt32 => {
+            let ca: &mut ChunkedArray<UInt32Type> = mutable_s.as_mut();
+            let values = values.u32()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::UInt64 => {
+            let ca: &mut ChunkedArray<UInt64Type> = mutable_s.as_mut();
+            let values = values.u64()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::Float32 => {
+            let ca: &mut ChunkedArray<Float32Type> = mutable_s.as_mut();
+            let values = values.f32()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::Float64 => {
+            let ca: &mut ChunkedArray<Float64Type> = mutable_s.as_mut();
+            let values = values.f64()?;
+            std::mem::take(ca).set_at_idx2(idx, values.into_iter())
+        }
+        DataType::Boolean => {
+            let ca = s.bool()?;
+            let values = values.bool()?;
+            ca.set_at_idx2(idx, values)
+        }
+        DataType::Utf8 => {
+            let ca = s.utf8()?;
+            let values = values.utf8()?;
+            ca.set_at_idx2(idx, values)
+        }
+        _ => panic!("not yet implemented for dtype: {}", logical_dtype),
+    };
+
+    s.and_then(|s| s.cast(&logical_dtype))
+}

data/ext/polars/src/utils.rs

@@ -0,0 +1,19 @@
+use polars::prelude::*;
+
+pub fn reinterpret(s: &Series, signed: bool) -> polars::prelude::PolarsResult<Series> {
+    match (s.dtype(), signed) {
+        (DataType::UInt64, true) => {
+            let ca = s.u64().unwrap();
+            Ok(ca.reinterpret_signed().into_series())
+        }
+        (DataType::UInt64, false) => Ok(s.clone()),
+        (DataType::Int64, false) => {
+            let ca = s.i64().unwrap();
+            Ok(ca.reinterpret_unsigned().into_series())
+        }
+        (DataType::Int64, true) => Ok(s.clone()),
+        _ => Err(PolarsError::ComputeError(
+            "reinterpret is only allowed for 64bit integers dtype, use cast otherwise".into(),
+        )),
+    }
+}

data/lib/polars/batched_csv_reader.rb

@@ -0,0 +1,96 @@
+module Polars
+  # @private
+  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

@@ -1,11 +1,50 @@
 module Polars
+  # Namespace for categorical related expressions.
   class CatExpr
+    # @private
     attr_accessor :_rbexpr
 
+    # @private
     def initialize(expr)
       self._rbexpr = expr._rbexpr
     end
 
+    # Determine how this categorical series should be sorted.
+    #
+    # @param ordering ["physical", "lexical"]
+    #   Ordering type:
+    #
+    #   - 'physical' -> Use the physical representation of the categories to determine the order (default).
+    #   - 'lexical' -> Use the string values to determine the ordering.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"cats" => ["z", "z", "k", "a", "b"], "vals" => [3, 1, 2, 2, 3]}
+    #   ).with_columns(
+    #     [
+    #       Polars.col("cats").cast(:cat).cat.set_ordering("lexical")
+    #     ]
+    #   )
+    #   df.sort(["cats", "vals"])
+    #   # =>
+    #   # shape: (5, 2)
+    #   # ┌──────┬──────┐
+    #   # │ cats ┆ vals │
+    #   # │ ---  ┆ ---  │
+    #   # │ cat  ┆ i64  │
+    #   # ╞══════╪══════╡
+    #   # │ a    ┆ 2    │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌┤
+    #   # │ b    ┆ 3    │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌┤
+    #   # │ k    ┆ 2    │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌┤
+    #   # │ z    ┆ 1    │
+    #   # ├╌╌╌╌╌╌┼╌╌╌╌╌╌┤
+    #   # │ z    ┆ 3    │
+    #   # └──────┴──────┘
     def set_ordering(ordering)
       Utils.wrap_expr(_rbexpr.cat_set_ordering(ordering))
     end