polars-df 0.1.4 → 0.2.0

data/ext/polars/src/series.rs

@@ -4,6 +4,8 @@ use polars::prelude::*;
 use polars::series::IsSorted;
 use std::cell::RefCell;
 
+use crate::apply::series::{call_lambda_and_extract, ApplyLambda};
+use crate::apply_method_all_arrow_series2;
 use crate::conversion::*;
 use crate::list_construction::rb_seq_to_list;
 use crate::set::set_at_idx;
@@ -142,7 +144,7 @@ impl RbSeries {
     }
 
     pub fn get_fmt(&self, index: usize, str_lengths: usize) -> String {
-        let val = format!("{}", self.series.borrow().get(index));
+        let val = format!("{}", self.series.borrow().get(index).unwrap());
         if let DataType::Utf8 | DataType::Categorical(_) = self.series.borrow().dtype() {
             let v_trunc = &val[..val
                 .char_indices()
@@ -170,8 +172,8 @@ impl RbSeries {
         }
     }
 
-    pub fn get_idx(&self, idx: usize) -> Value {
-        Wrap(self.series.borrow().get(idx)).into()
+    pub fn get_idx(&self, idx: usize) -> RbResult<Value> {
+        Ok(Wrap(self.series.borrow().get(idx).map_err(RbPolarsErr::from)?).into())
     }
 
     pub fn bitand(&self, other: &RbSeries) -> RbResult<Self> {
@@ -245,16 +247,37 @@ impl RbSeries {
         }
     }
 
-    pub fn max(&self) -> Value {
-        Wrap(self.series.borrow().max_as_series().get(0)).into()
+    pub fn max(&self) -> RbResult<Value> {
+        Ok(Wrap(
+            self.series
+                .borrow()
+                .max_as_series()
+                .get(0)
+                .map_err(RbPolarsErr::from)?,
+        )
+        .into())
     }
 
-    pub fn min(&self) -> Value {
-        Wrap(self.series.borrow().min_as_series().get(0)).into()
+    pub fn min(&self) -> RbResult<Value> {
+        Ok(Wrap(
+            self.series
+                .borrow()
+                .min_as_series()
+                .get(0)
+                .map_err(RbPolarsErr::from)?,
+        )
+        .into())
     }
 
-    pub fn sum(&self) -> Value {
-        Wrap(self.series.borrow().sum_as_series().get(0)).into()
+    pub fn sum(&self) -> RbResult<Value> {
+        Ok(Wrap(
+            self.series
+                .borrow()
+                .sum_as_series()
+                .get(0)
+                .map_err(RbPolarsErr::from)?,
+        )
+        .into())
     }
 
     pub fn n_chunks(&self) -> usize {
@@ -520,7 +543,8 @@ impl RbSeries {
                 .borrow()
                 .quantile_as_series(quantile, interpolation.0)
                 .map_err(|_| RbValueError::new_err("invalid quantile".into()))?
-                .get(0),
+                .get(0)
+                .unwrap_or(AnyValue::Null),
         )
         .into())
     }
@@ -529,6 +553,198 @@ impl RbSeries {
         RbSeries::new(self.series.borrow().clone())
     }
 
+    pub fn apply_lambda(
+        &self,
+        lambda: Value,
+        output_type: Option<Wrap<DataType>>,
+        skip_nulls: bool,
+    ) -> RbResult<Self> {
+        let series = &self.series.borrow();
+
+        let output_type = output_type.map(|dt| dt.0);
+
+        macro_rules! dispatch_apply {
+            ($self:expr, $method:ident, $($args:expr),*) => {
+                if matches!($self.dtype(), DataType::Object(_)) {
+                    // let ca = $self.0.unpack::<ObjectType<ObjectValue>>().unwrap();
+                    // ca.$method($($args),*)
+                    todo!()
+                } else {
+                    apply_method_all_arrow_series2!(
+                        $self,
+                        $method,
+                        $($args),*
+                    )
+                }
+
+            }
+
+        }
+
+        if matches!(
+            series.dtype(),
+            DataType::Datetime(_, _)
+                | DataType::Date
+                | DataType::Duration(_)
+                | DataType::Categorical(_)
+                | DataType::Time
+        ) || !skip_nulls
+        {
+            let mut avs = Vec::with_capacity(series.len());
+            let iter = series.iter().map(|av| {
+                let input = Wrap(av);
+                call_lambda_and_extract::<_, Wrap<AnyValue>>(lambda, input)
+                    .unwrap()
+                    .0
+            });
+            avs.extend(iter);
+            return Ok(Series::new(&self.name(), &avs).into());
+        }
+
+        let out = match output_type {
+            Some(DataType::Int8) => {
+                let ca: Int8Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::Int16) => {
+                let ca: Int16Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::Int32) => {
+                let ca: Int32Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::Int64) => {
+                let ca: Int64Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::UInt8) => {
+                let ca: UInt8Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::UInt16) => {
+                let ca: UInt16Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::UInt32) => {
+                let ca: UInt32Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::UInt64) => {
+                let ca: UInt64Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::Float32) => {
+                let ca: Float32Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::Float64) => {
+                let ca: Float64Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_series()
+            }
+            Some(DataType::Boolean) => {
+                let ca: BooleanChunked =
+                    dispatch_apply!(series, apply_lambda_with_bool_out_type, lambda, 0, None)?;
+                ca.into_series()
+            }
+            Some(DataType::Date) => {
+                let ca: Int32Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_date().into_series()
+            }
+            Some(DataType::Datetime(tu, tz)) => {
+                let ca: Int64Chunked = dispatch_apply!(
+                    series,
+                    apply_lambda_with_primitive_out_type,
+                    lambda,
+                    0,
+                    None
+                )?;
+                ca.into_datetime(tu, tz).into_series()
+            }
+            Some(DataType::Utf8) => {
+                let ca = dispatch_apply!(series, apply_lambda_with_utf8_out_type, lambda, 0, None)?;
+
+                ca.into_series()
+            }
+            Some(DataType::Object(_)) => {
+                let ca =
+                    dispatch_apply!(series, apply_lambda_with_object_out_type, lambda, 0, None)?;
+                ca.into_series()
+            }
+            None => return dispatch_apply!(series, apply_lambda_unknown, lambda),
+
+            _ => return dispatch_apply!(series, apply_lambda_unknown, lambda),
+        };
+
+        Ok(RbSeries::new(out))
+    }
+
     pub fn zip_with(&self, mask: &RbSeries, other: &RbSeries) -> RbResult<Self> {
         let binding = mask.series.borrow();
         let mask = binding.bool().map_err(RbPolarsErr::from)?;
@@ -627,6 +843,108 @@ impl RbSeries {
     }
 }
 
+macro_rules! impl_set_with_mask {
+    ($name:ident, $native:ty, $cast:ident, $variant:ident) => {
+        fn $name(
+            series: &Series,
+            filter: &RbSeries,
+            value: Option<$native>,
+        ) -> PolarsResult<Series> {
+            let binding = filter.series.borrow();
+            let mask = binding.bool()?;
+            let ca = series.$cast()?;
+            let new = ca.set(mask, value)?;
+            Ok(new.into_series())
+        }
+
+        impl RbSeries {
+            pub fn $name(&self, filter: &RbSeries, value: Option<$native>) -> RbResult<Self> {
+                let series =
+                    $name(&self.series.borrow(), filter, value).map_err(RbPolarsErr::from)?;
+                Ok(Self::new(series))
+            }
+        }
+    };
+}
+
+// impl_set_with_mask!(set_with_mask_str, &str, utf8, Utf8);
+impl_set_with_mask!(set_with_mask_f64, f64, f64, Float64);
+impl_set_with_mask!(set_with_mask_f32, f32, f32, Float32);
+impl_set_with_mask!(set_with_mask_u8, u8, u8, UInt8);
+impl_set_with_mask!(set_with_mask_u16, u16, u16, UInt16);
+impl_set_with_mask!(set_with_mask_u32, u32, u32, UInt32);
+impl_set_with_mask!(set_with_mask_u64, u64, u64, UInt64);
+impl_set_with_mask!(set_with_mask_i8, i8, i8, Int8);
+impl_set_with_mask!(set_with_mask_i16, i16, i16, Int16);
+impl_set_with_mask!(set_with_mask_i32, i32, i32, Int32);
+impl_set_with_mask!(set_with_mask_i64, i64, i64, Int64);
+impl_set_with_mask!(set_with_mask_bool, bool, bool, Boolean);
+
+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, %);
+
 macro_rules! impl_eq_num {
     ($name:ident, $type:ty) => {
         impl RbSeries {

data/ext/polars/src/utils.rs

@@ -17,3 +17,28 @@ pub fn reinterpret(s: &Series, signed: bool) -> polars::prelude::PolarsResult<Se
         )),
     }
 }
+
+#[macro_export]
+macro_rules! apply_method_all_arrow_series2 {
+    ($self:expr, $method:ident, $($args:expr),*) => {
+        match $self.dtype() {
+            DataType::Boolean => $self.bool().unwrap().$method($($args),*),
+            DataType::Utf8 => $self.utf8().unwrap().$method($($args),*),
+            DataType::UInt8 => $self.u8().unwrap().$method($($args),*),
+            DataType::UInt16 => $self.u16().unwrap().$method($($args),*),
+            DataType::UInt32 => $self.u32().unwrap().$method($($args),*),
+            DataType::UInt64 => $self.u64().unwrap().$method($($args),*),
+            DataType::Int8 => $self.i8().unwrap().$method($($args),*),
+            DataType::Int16 => $self.i16().unwrap().$method($($args),*),
+            DataType::Int32 => $self.i32().unwrap().$method($($args),*),
+            DataType::Int64 => $self.i64().unwrap().$method($($args),*),
+            DataType::Float32 => $self.f32().unwrap().$method($($args),*),
+            DataType::Float64 => $self.f64().unwrap().$method($($args),*),
+            DataType::Date => $self.date().unwrap().$method($($args),*),
+            DataType::Datetime(_, _) => $self.datetime().unwrap().$method($($args),*),
+            // DataType::List(_) => $self.list().unwrap().$method($($args),*),
+            DataType::Struct(_) => $self.struct_().unwrap().$method($($args),*),
+            dt => panic!("dtype {:?} not supported", dt)
+        }
+    }
+}

data/lib/polars/convert.rb

@@ -0,0 +1,100 @@
+module Polars
+  module Convert
+    # Construct a DataFrame from a dictionary of sequences.
+    #
+    # This operation clones data, unless you pass in a `Hash<String, Series>`.
+    #
+    # @param data [Hash]
+    #   Two-dimensional data represented as a hash. Hash must contain
+    #   arrays.
+    # @param columns [Array]
+    #   Column labels to use for resulting DataFrame. If specified, overrides any
+    #   labels already present in the data. Must match data dimensions.
+    #
+    # @return [DataFrame]
+    #
+    # @example
+    #   data = {"a" => [1, 2], "b" => [3, 4]}
+    #   Polars.from_hash(data)
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 3   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ 2   ┆ 4   │
+    #   # └─────┴─────┘
+    def from_hash(data, columns: nil)
+      DataFrame._from_hash(data, columns: columns)
+    end
+
+    # Construct a DataFrame from a sequence of dictionaries. This operation clones data.
+    #
+    # @param hashes [Array]
+    #   Array with hashes mapping column name to value.
+    # @param infer_schema_length [Integer]
+    #   How many hashes/rows to scan to determine the data types
+    #   if set to `nil` all rows are scanned. This will be slow.
+    # @param schema [Object]
+    #   Schema that (partially) overwrites the inferred schema.
+    #
+    # @return [DataFrame]
+    #
+    # @example
+    #   data = [{"a" => 1, "b" => 4}, {"a" => 2, "b" => 5}, {"a" => 3, "b" => 6}]
+    #   Polars.from_hashes(data)
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 4   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ 2   ┆ 5   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ 3   ┆ 6   │
+    #   # └─────┴─────┘
+    #
+    # @example Overwrite first column name and dtype
+    #   Polars.from_hashes(data, schema: {"c" => :i32})
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ c   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i32 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1   ┆ 4   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ 2   ┆ 5   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ 3   ┆ 6   │
+    #   # └─────┴─────┘
+    #
+    # @example Let polars infer the dtypes but inform about a 3rd column
+    #   Polars.from_hashes(data, schema: {"a" => :unknown, "b" => :unknown, "c" => :i32})
+    #   # shape: (3, 3)
+    #   # ┌─────┬─────┬──────┐
+    #   # │ a   ┆ b   ┆ c    │
+    #   # │ --- ┆ --- ┆ ---  │
+    #   # │ i64 ┆ i64 ┆ i32  │
+    #   # ╞═════╪═════╪══════╡
+    #   # │ 1   ┆ 4   ┆ null │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┤
+    #   # │ 2   ┆ 5   ┆ null │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌┤
+    #   # │ 3   ┆ 6   ┆ null │
+    #   # └─────┴─────┴──────┘
+    # def from_hashes(hashes, infer_schema_length: 50, schema: nil)
+    #   DataFrame._from_hashes(hashes, infer_schema_length: infer_schema_length, schema: schema)
+    # end
+
+    # def from_records
+    # end
+  end
+end