polars-df 0.1.3 → 0.1.5

data/ext/polars/src/series.rs

@@ -4,7 +4,10 @@ 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;
 use crate::{RbDataFrame, RbPolarsErr, RbResult, RbValueError};
 
@@ -123,6 +126,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()
     }
@@ -515,6 +531,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)?;
@@ -613,6 +821,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 {
@@ -787,4 +1097,25 @@ impl RbSeries {
         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/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/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

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