polars-df 0.4.0 → 0.5.0

data/ext/polars/src/rb_modules.rs

@@ -8,6 +8,14 @@ pub(crate) fn series() -> RClass {
     *memoize!(RClass: polars().const_get("Series").unwrap())
 }
 
+pub(crate) fn utils() -> RModule {
+    *memoize!(RModule: polars().const_get("Utils").unwrap())
+}
+
 pub(crate) fn date() -> RClass {
     *memoize!(RClass: class::object().const_get("Date").unwrap())
 }
+
+pub(crate) fn datetime() -> RClass {
+    *memoize!(RClass: class::object().const_get("DateTime").unwrap())
+}

data/ext/polars/src/series/aggregation.rs

@@ -0,0 +1,83 @@
+use crate::error::RbPolarsErr;
+use crate::prelude::*;
+use crate::{RbResult, RbSeries, RbValueError};
+use magnus::{IntoValue, Value};
+
+impl RbSeries {
+    pub fn arg_max(&self) -> Option<usize> {
+        self.series.borrow().arg_max()
+    }
+
+    pub fn arg_min(&self) -> Option<usize> {
+        self.series.borrow().arg_min()
+    }
+
+    pub fn max(&self) -> RbResult<Value> {
+        Ok(Wrap(
+            self.series
+                .borrow()
+                .max_as_series()
+                .get(0)
+                .map_err(RbPolarsErr::from)?,
+        )
+        .into_value())
+    }
+
+    pub fn mean(&self) -> Option<f64> {
+        match self.series.borrow().dtype() {
+            DataType::Boolean => {
+                let s = self.series.borrow().cast(&DataType::UInt8).unwrap();
+                s.mean()
+            }
+            _ => self.series.borrow().mean(),
+        }
+    }
+
+    pub fn median(&self) -> Option<f64> {
+        match self.series.borrow().dtype() {
+            DataType::Boolean => {
+                let s = self.series.borrow().cast(&DataType::UInt8).unwrap();
+                s.median()
+            }
+            _ => self.series.borrow().median(),
+        }
+    }
+
+    pub fn min(&self) -> RbResult<Value> {
+        Ok(Wrap(
+            self.series
+                .borrow()
+                .min_as_series()
+                .get(0)
+                .map_err(RbPolarsErr::from)?,
+        )
+        .into_value())
+    }
+
+    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)
+                .unwrap_or(AnyValue::Null),
+        )
+        .into_value())
+    }
+
+    pub fn sum(&self) -> RbResult<Value> {
+        Ok(Wrap(
+            self.series
+                .borrow()
+                .sum_as_series()
+                .get(0)
+                .map_err(RbPolarsErr::from)?,
+        )
+        .into_value())
+    }
+}

data/ext/polars/src/series/arithmetic.rs

@@ -0,0 +1,88 @@
+use crate::{RbResult, RbSeries};
+
+impl RbSeries {
+    pub fn add(&self, other: &RbSeries) -> Self {
+        (&*self.series.borrow() + &*other.series.borrow()).into()
+    }
+
+    pub fn sub(&self, other: &RbSeries) -> Self {
+        (&*self.series.borrow() - &*other.series.borrow()).into()
+    }
+
+    pub fn mul(&self, other: &RbSeries) -> Self {
+        (&*self.series.borrow() * &*other.series.borrow()).into()
+    }
+
+    pub fn div(&self, other: &RbSeries) -> Self {
+        (&*self.series.borrow() / &*other.series.borrow()).into()
+    }
+
+    pub fn rem(&self, other: &RbSeries) -> Self {
+        (&*self.series.borrow() % &*other.series.borrow()).into()
+    }
+}
+
+macro_rules! impl_arithmetic {
+    ($name:ident, $type:ty, $operand:tt) => {
+        impl RbSeries {
+            pub fn $name(&self, other: $type) -> RbResult<Self> {
+                Ok(RbSeries::new(&*self.series.borrow() $operand other))
+            }
+        }
+    };
+}
+
+impl_arithmetic!(add_u8, u8, +);
+impl_arithmetic!(add_u16, u16, +);
+impl_arithmetic!(add_u32, u32, +);
+impl_arithmetic!(add_u64, u64, +);
+impl_arithmetic!(add_i8, i8, +);
+impl_arithmetic!(add_i16, i16, +);
+impl_arithmetic!(add_i32, i32, +);
+impl_arithmetic!(add_i64, i64, +);
+impl_arithmetic!(add_datetime, i64, +);
+impl_arithmetic!(add_duration, i64, +);
+impl_arithmetic!(add_f32, f32, +);
+impl_arithmetic!(add_f64, f64, +);
+impl_arithmetic!(sub_u8, u8, -);
+impl_arithmetic!(sub_u16, u16, -);
+impl_arithmetic!(sub_u32, u32, -);
+impl_arithmetic!(sub_u64, u64, -);
+impl_arithmetic!(sub_i8, i8, -);
+impl_arithmetic!(sub_i16, i16, -);
+impl_arithmetic!(sub_i32, i32, -);
+impl_arithmetic!(sub_i64, i64, -);
+impl_arithmetic!(sub_datetime, i64, -);
+impl_arithmetic!(sub_duration, i64, -);
+impl_arithmetic!(sub_f32, f32, -);
+impl_arithmetic!(sub_f64, f64, -);
+impl_arithmetic!(div_u8, u8, /);
+impl_arithmetic!(div_u16, u16, /);
+impl_arithmetic!(div_u32, u32, /);
+impl_arithmetic!(div_u64, u64, /);
+impl_arithmetic!(div_i8, i8, /);
+impl_arithmetic!(div_i16, i16, /);
+impl_arithmetic!(div_i32, i32, /);
+impl_arithmetic!(div_i64, i64, /);
+impl_arithmetic!(div_f32, f32, /);
+impl_arithmetic!(div_f64, f64, /);
+impl_arithmetic!(mul_u8, u8, *);
+impl_arithmetic!(mul_u16, u16, *);
+impl_arithmetic!(mul_u32, u32, *);
+impl_arithmetic!(mul_u64, u64, *);
+impl_arithmetic!(mul_i8, i8, *);
+impl_arithmetic!(mul_i16, i16, *);
+impl_arithmetic!(mul_i32, i32, *);
+impl_arithmetic!(mul_i64, i64, *);
+impl_arithmetic!(mul_f32, f32, *);
+impl_arithmetic!(mul_f64, f64, *);
+impl_arithmetic!(rem_u8, u8, %);
+impl_arithmetic!(rem_u16, u16, %);
+impl_arithmetic!(rem_u32, u32, %);
+impl_arithmetic!(rem_u64, u64, %);
+impl_arithmetic!(rem_i8, i8, %);
+impl_arithmetic!(rem_i16, i16, %);
+impl_arithmetic!(rem_i32, i32, %);
+impl_arithmetic!(rem_i64, i64, %);
+impl_arithmetic!(rem_f32, f32, %);
+impl_arithmetic!(rem_f64, f64, %);

data/ext/polars/src/series/comparison.rs

@@ -0,0 +1,251 @@
+use crate::error::RbPolarsErr;
+use crate::prelude::*;
+use crate::{RbResult, RbSeries};
+
+impl RbSeries {
+    pub fn eq(&self, rhs: &RbSeries) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .equal(&*rhs.series.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(Self::new(s.into_series()))
+    }
+
+    pub fn neq(&self, rhs: &RbSeries) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .not_equal(&*rhs.series.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(Self::new(s.into_series()))
+    }
+
+    pub fn gt(&self, rhs: &RbSeries) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .gt(&*rhs.series.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(Self::new(s.into_series()))
+    }
+
+    pub fn gt_eq(&self, rhs: &RbSeries) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .gt_eq(&*rhs.series.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(Self::new(s.into_series()))
+    }
+
+    pub fn lt(&self, rhs: &RbSeries) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .lt(&*rhs.series.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(Self::new(s.into_series()))
+    }
+
+    pub fn lt_eq(&self, rhs: &RbSeries) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .lt_eq(&*rhs.series.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(Self::new(s.into_series()))
+    }
+}
+
+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()))
+            }
+        }
+    };
+}
+
+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);
+
+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()))
+            }
+        }
+    };
+}
+
+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);
+
+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()))
+            }
+        }
+    };
+}
+
+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);
+
+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()))
+            }
+        }
+    };
+}
+
+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);
+
+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()))
+            }
+        }
+    };
+}
+
+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);
+
+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 RbSeries {
+    pub fn eq_str(&self, rhs: String) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .equal(rhs.as_str())
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbSeries::new(s.into_series()))
+    }
+
+    pub fn neq_str(&self, rhs: String) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .not_equal(rhs.as_str())
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbSeries::new(s.into_series()))
+    }
+
+    pub fn gt_str(&self, rhs: String) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .gt(rhs.as_str())
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbSeries::new(s.into_series()))
+    }
+
+    pub fn gt_eq_str(&self, rhs: String) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .gt_eq(rhs.as_str())
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbSeries::new(s.into_series()))
+    }
+
+    pub fn lt_str(&self, rhs: String) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .lt(rhs.as_str())
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbSeries::new(s.into_series()))
+    }
+
+    pub fn lt_eq_str(&self, rhs: String) -> RbResult<Self> {
+        let s = self
+            .series
+            .borrow()
+            .lt_eq(rhs.as_str())
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbSeries::new(s.into_series()))
+    }
+}

data/ext/polars/src/series/construction.rs

@@ -0,0 +1,164 @@
+use magnus::RArray;
+use polars_core::prelude::*;
+
+use crate::conversion::{slice_extract_wrapped, Wrap};
+use crate::prelude::ObjectValue;
+use crate::series::to_series_collection;
+use crate::{RbPolarsErr, RbResult, RbSeries};
+
+impl RbSeries {
+    pub fn new_opt_bool(name: String, obj: RArray, strict: bool) -> RbResult<RbSeries> {
+        let len = obj.len();
+        let mut builder = BooleanChunkedBuilder::new(&name, len);
+
+        unsafe {
+            for item in obj.as_slice().iter() {
+                if item.is_nil() {
+                    builder.append_null()
+                } else {
+                    match item.try_convert::<bool>() {
+                        Ok(val) => builder.append_value(val),
+                        Err(e) => {
+                            if strict {
+                                return Err(e);
+                            }
+                            builder.append_null()
+                        }
+                    }
+                }
+            }
+        }
+        let ca = builder.finish();
+
+        let s = ca.into_series();
+        Ok(RbSeries::new(s))
+    }
+}
+
+fn new_primitive<T>(name: &str, obj: RArray, strict: bool) -> RbResult<RbSeries>
+where
+    T: PolarsNumericType,
+    ChunkedArray<T>: IntoSeries,
+    T::Native: magnus::TryConvert,
+{
+    let len = obj.len();
+    let mut builder = PrimitiveChunkedBuilder::<T>::new(name, len);
+
+    unsafe {
+        for item in obj.as_slice().iter() {
+            if item.is_nil() {
+                builder.append_null()
+            } else {
+                match item.try_convert::<T::Native>() {
+                    Ok(val) => builder.append_value(val),
+                    Err(e) => {
+                        if strict {
+                            return Err(e);
+                        }
+                        builder.append_null()
+                    }
+                }
+            }
+        }
+    }
+    let ca = builder.finish();
+
+    let s = ca.into_series();
+    Ok(RbSeries::new(s))
+}
+
+// Init with lists that can contain Nones
+macro_rules! init_method_opt {
+    ($name:ident, $type:ty, $native: ty) => {
+        impl RbSeries {
+            pub fn $name(name: String, obj: RArray, strict: bool) -> RbResult<Self> {
+                new_primitive::<$type>(&name, obj, strict)
+            }
+        }
+    };
+}
+
+init_method_opt!(new_opt_u8, UInt8Type, u8);
+init_method_opt!(new_opt_u16, UInt16Type, u16);
+init_method_opt!(new_opt_u32, UInt32Type, u32);
+init_method_opt!(new_opt_u64, UInt64Type, u64);
+init_method_opt!(new_opt_i8, Int8Type, i8);
+init_method_opt!(new_opt_i16, Int16Type, i16);
+init_method_opt!(new_opt_i32, Int32Type, i32);
+init_method_opt!(new_opt_i64, Int64Type, i64);
+init_method_opt!(new_opt_f32, Float32Type, f32);
+init_method_opt!(new_opt_f64, Float64Type, f64);
+
+fn vec_wrap_any_value<'s>(arr: RArray) -> RbResult<Vec<Wrap<AnyValue<'s>>>> {
+    let mut val = Vec::with_capacity(arr.len());
+    for v in arr.each() {
+        val.push(v?.try_convert()?);
+    }
+    Ok(val)
+}
+
+impl RbSeries {
+    pub fn new_from_anyvalues(name: String, val: RArray, strict: bool) -> RbResult<Self> {
+        let val = vec_wrap_any_value(val)?;
+        let avs = slice_extract_wrapped(&val);
+        // from anyvalues is fallible
+        let s = Series::from_any_values(&name, avs, strict).map_err(RbPolarsErr::from)?;
+        Ok(s.into())
+    }
+
+    pub fn new_str(name: String, val: Wrap<Utf8Chunked>, _strict: bool) -> Self {
+        let mut s = val.0.into_series();
+        s.rename(&name);
+        RbSeries::new(s)
+    }
+
+    pub fn new_binary(name: String, val: Wrap<BinaryChunked>, _strict: bool) -> Self {
+        let mut s = val.0.into_series();
+        s.rename(&name);
+        RbSeries::new(s)
+    }
+
+    pub fn new_null(name: String, val: RArray, _strict: bool) -> RbResult<Self> {
+        let s = Series::new_null(&name, val.len());
+        Ok(s.into())
+    }
+
+    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_series_list(name: String, val: RArray, _strict: bool) -> RbResult<Self> {
+        let series_vec = to_series_collection(val)?;
+        Ok(Series::new(&name, &series_vec).into())
+    }
+
+    pub fn new_decimal(name: String, val: RArray, strict: bool) -> RbResult<Self> {
+        let val = vec_wrap_any_value(val)?;
+        // TODO: do we have to respect 'strict' here? it's possible if we want to
+        let avs = slice_extract_wrapped(&val);
+        // create a fake dtype with a placeholder "none" scale, to be inferred later
+        let dtype = DataType::Decimal(None, None);
+        let s = Series::from_any_values_and_dtype(&name, avs, &dtype, strict)
+            .map_err(RbPolarsErr::from)?;
+        Ok(s.into())
+    }
+
+    pub fn repeat(
+        name: String,
+        val: Wrap<AnyValue>,
+        n: usize,
+        dtype: Wrap<DataType>,
+    ) -> RbResult<Self> {
+        let av = val.0;
+        Ok(Series::new(&name, &[av])
+            .cast(&dtype.0)
+            .map_err(RbPolarsErr::from)?
+            .new_from_index(0, n)
+            .into())
+    }
+}