polars-df 0.1.2 → 0.1.4

data/ext/polars/src/series.rs

@@ -5,6 +5,8 @@ use polars::series::IsSorted;
 use std::cell::RefCell;
 
 use crate::conversion::*;
+use crate::list_construction::rb_seq_to_list;
+use crate::set::set_at_idx;
 use crate::{RbDataFrame, RbPolarsErr, RbResult, RbValueError};
 
 #[magnus::wrap(class = "Polars::RbSeries")]
@@ -122,6 +124,19 @@ impl RbSeries {
         RbSeries::new(s)
     }
 
+    pub fn new_object(name: String, val: RArray, _strict: bool) -> RbResult<Self> {
+        let val = val
+            .each()
+            .map(|v| v.map(ObjectValue::from))
+            .collect::<RbResult<Vec<ObjectValue>>>()?;
+        let s = ObjectChunked::<ObjectValue>::new_from_vec(&name, val).into_series();
+        Ok(s.into())
+    }
+
+    pub fn new_list(name: String, seq: Value, dtype: Wrap<DataType>) -> RbResult<Self> {
+        rb_seq_to_list(&name, seq, &dtype.0).map(|s| s.into())
+    }
+
     pub fn estimated_size(&self) -> usize {
         self.series.borrow().estimated_size()
     }
@@ -474,6 +489,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!();
         }
@@ -594,44 +615,211 @@ 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(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())
+    }
+
+    pub fn new_opt_datetime(name: String, values: RArray, _strict: Option<bool>) -> RbResult<Self> {
+        let len = values.len();
+        let mut builder = PrimitiveChunkedBuilder::<Int64Type>::new(&name, len);
+        for item in values.each() {
+            let v = item?;
+            if v.is_nil() {
+                builder.append_null();
+            } else {
+                let sec: i64 = v.funcall("to_i", ())?;
+                let nsec: i64 = v.funcall("nsec", ())?;
+                // TODO use strict
+                builder.append_value(sec * 1_000_000_000 + nsec);
+            }
+        }
+        let ca: ChunkedArray<Int64Type> = builder.finish();
+        Ok(ca
+            .into_datetime(TimeUnit::Nanoseconds, None)
+            .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

@@ -1,4 +1,5 @@
 module Polars
+  # @private
   class BatchedCsvReader
     attr_accessor :_reader, :new_columns
 

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

data/lib/polars/cat_name_space.rb

@@ -0,0 +1,54 @@
+module Polars
+  # Series.cat namespace.
+  class CatNameSpace
+    include ExprDispatch
+
+    self._accessor = "cat"
+
+    # @private
+    def initialize(series)
+      self._s = series._s
+    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 [Series]
+    #
+    # @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)
+      super
+    end
+  end
+end