polars-df 0.5.0 → 0.6.0

data/README.md

@@ -25,7 +25,7 @@ Polars.read_csv("iris.csv")
   .collect
 ```
 
-You can follow [Polars tutorials](https://pola-rs.github.io/polars-book/user-guide/introduction.html) and convert the code to Ruby in many cases. Feel free to open an issue if you run into problems.
+You can follow [Polars tutorials](https://pola-rs.github.io/polars-book/user-guide/) and convert the code to Ruby in many cases. Feel free to open an issue if you run into problems.
 
 ## Reference
 
@@ -348,7 +348,7 @@ df.to_numo
 You can specify column types when creating a data frame
 
 ```ruby
-Polars::DataFrame.new(data, columns: {"a" => Polars::Int32, "b" => Polars::Float32})
+Polars::DataFrame.new(data, schema: {"a" => Polars::Int32, "b" => Polars::Float32})
 ```
 
 Supported types are:
@@ -357,8 +357,9 @@ Supported types are:
 - float - `Float64`, `Float32`
 - integer - `Int64`, `Int32`, `Int16`, `Int8`
 - unsigned integer - `UInt64`, `UInt32`, `UInt16`, `UInt8`
-- string - `Utf8`, `Categorical`
+- string - `Utf8`, `Binary`, `Categorical`
 - temporal - `Date`, `Datetime`, `Time`, `Duration`
+- other - `Object`, `List`, `Struct`, `Array` [unreleased]
 
 Get column types
 

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.5.0"
+version = "0.6.0"
 license = "MIT"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
@@ -11,17 +11,17 @@ crate-type = ["cdylib"]
 
 [dependencies]
 ahash = "0.8"
+chrono = "=0.4.24"
 magnus = "0.5"
-polars-core = "0.29.0"
+polars-core = "0.31.1"
 serde_json = "1"
 smartstring = "1"
 
 [dependencies.polars]
-version = "0.29.0"
+version = "0.31.1"
 features = [
     "abs",
     "approx_unique",
-    "arange",
     "arg_where",
     "asof_join",
     "avro",
@@ -68,6 +68,7 @@ features = [
     "propagate_nans",
     "random",
     "rank",
+    "range",
     "reinterpret",
     "repeat_by",
     "rolling_window",

data/ext/polars/src/apply/mod.rs

@@ -237,15 +237,19 @@ fn iterator_to_list(
     for _ in 0..init_null_count {
         builder.append_null()
     }
-    builder.append_opt_series(first_value);
+    builder
+        .append_opt_series(first_value)
+        .map_err(RbPolarsErr::from)?;
     for opt_val in it {
         match opt_val {
             None => builder.append_null(),
             Some(s) => {
                 if s.len() == 0 && s.dtype() != dt {
-                    builder.append_series(&Series::full_null("", 0, dt))
+                    builder
+                        .append_series(&Series::full_null("", 0, dt))
+                        .unwrap()
                 } else {
-                    builder.append_series(&s)
+                    builder.append_series(&s).map_err(RbPolarsErr::from)?
                 }
             }
         }

data/ext/polars/src/conversion.rs

@@ -1,11 +1,11 @@
 use std::fmt::{Display, Formatter};
 use std::hash::{Hash, Hasher};
 
+use magnus::encoding::{EncodingCapable, Index};
 use magnus::{
-    class, exception, r_hash::ForEach, ruby_handle::RubyHandle, Integer, IntoValue, Module, RArray,
-    RFloat, RHash, RString, Symbol, TryConvert, Value, QNIL,
+    class, exception, r_hash::ForEach, ruby_handle::RubyHandle, Float, Integer, IntoValue, Module,
+    RArray, RHash, RString, Symbol, TryConvert, Value, QNIL,
 };
-use magnus::encoding::{EncodingCapable, Index};
 use polars::chunked_array::object::PolarsObjectSafe;
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
 use polars::datatypes::AnyValue;
@@ -16,8 +16,9 @@ use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 use smartstring::alias::String as SmartString;
 
+use crate::object::OBJECT_NAME;
 use crate::rb_modules::utils;
-use crate::{RbDataFrame, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
+use crate::{RbDataFrame, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbTypeError, RbValueError};
 
 pub(crate) fn slice_to_wrapped<T>(slice: &[T]) -> &[Wrap<T>] {
     // Safety:
@@ -149,51 +150,29 @@ impl IntoValue for Wrap<AnyValue<'_>> {
             AnyValue::Boolean(v) => Value::from(v),
             AnyValue::Utf8(v) => Value::from(v),
             AnyValue::Utf8Owned(v) => Value::from(v.as_str()),
-            AnyValue::Categorical(_idx, _rev, _arr) => todo!(),
-            AnyValue::Date(v) => class::time()
-                .funcall::<_, _, Value>("at", (v * 86400,))
-                .unwrap()
-                .funcall::<_, _, Value>("utc", ())
-                .unwrap()
-                .funcall::<_, _, Value>("to_date", ())
-                .unwrap(),
-            AnyValue::Datetime(v, tu, tz) => {
-                let t = match tu {
-                    TimeUnit::Nanoseconds => {
-                        let sec = v / 1000000000;
-                        let subsec = v % 1000000000;
-                        class::time()
-                            .funcall::<_, _, Value>("at", (sec, subsec, Symbol::new("nsec")))
-                            .unwrap()
-                    }
-                    TimeUnit::Microseconds => {
-                        let sec = v / 1000000;
-                        let subsec = v % 1000000;
-                        class::time()
-                            .funcall::<_, _, Value>("at", (sec, subsec, Symbol::new("usec")))
-                            .unwrap()
-                    }
-                    TimeUnit::Milliseconds => {
-                        let sec = v / 1000;
-                        let subsec = v % 1000;
-                        class::time()
-                            .funcall::<_, _, Value>("at", (sec, subsec, Symbol::new("millisecond")))
-                            .unwrap()
-                    }
-                };
-
-                if tz.is_some() {
-                    todo!();
+            AnyValue::Categorical(idx, rev, arr) => {
+                let s = if arr.is_null() {
+                    rev.get(idx)
                 } else {
-                    t.funcall::<_, _, Value>("utc", ()).unwrap()
-                }
+                    unsafe { arr.deref_unchecked().value(idx as usize) }
+                };
+                s.into_value()
             }
-            AnyValue::Duration(v, tu) => {
-                let tu = tu.to_ascii();
-                utils().funcall("_to_ruby_duration", (v, tu)).unwrap()
+            AnyValue::Date(v) => utils().funcall("_to_ruby_date", (v,)).unwrap(),
+            AnyValue::Datetime(v, time_unit, time_zone) => {
+                let time_unit = time_unit.to_ascii();
+                utils()
+                    .funcall("_to_ruby_datetime", (v, time_unit, time_zone.clone()))
+                    .unwrap()
             }
-            AnyValue::Time(_v) => todo!(),
-            AnyValue::List(v) => RbSeries::new(v).to_a().into_value(),
+            AnyValue::Duration(v, time_unit) => {
+                let time_unit = time_unit.to_ascii();
+                utils()
+                    .funcall("_to_ruby_duration", (v, time_unit))
+                    .unwrap()
+            }
+            AnyValue::Time(v) => utils().funcall("_to_ruby_time", (v,)).unwrap(),
+            AnyValue::Array(v, _) | AnyValue::List(v) => RbSeries::new(v).to_a().into_value(),
             ref av @ AnyValue::Struct(_, _, flds) => struct_dict(av._iter_struct_av(), flds),
             AnyValue::StructOwned(payload) => struct_dict(payload.0.into_iter(), &payload.1),
             AnyValue::Object(v) => {
@@ -206,7 +185,9 @@ impl IntoValue for Wrap<AnyValue<'_>> {
             }
             AnyValue::Binary(v) => RString::from_slice(v).into_value(),
             AnyValue::BinaryOwned(v) => RString::from_slice(&v).into_value(),
-            AnyValue::Decimal(_v, _scale) => todo!(),
+            AnyValue::Decimal(v, scale) => utils()
+                .funcall("_to_ruby_decimal", (v.to_string(), -(scale as i32)))
+                .unwrap(),
         }
     }
 }
@@ -226,10 +207,22 @@ impl IntoValue for Wrap<DataType> {
             DataType::UInt64 => pl.const_get::<_, Value>("UInt64").unwrap(),
             DataType::Float32 => pl.const_get::<_, Value>("Float32").unwrap(),
             DataType::Float64 => pl.const_get::<_, Value>("Float64").unwrap(),
-            DataType::Decimal(_precision, _scale) => todo!(),
+            DataType::Decimal(precision, scale) => {
+                let decimal_class = pl.const_get::<_, Value>("Decimal").unwrap();
+                decimal_class
+                    .funcall::<_, _, Value>("new", (precision, scale))
+                    .unwrap()
+            }
             DataType::Boolean => pl.const_get::<_, Value>("Boolean").unwrap(),
             DataType::Utf8 => pl.const_get::<_, Value>("Utf8").unwrap(),
             DataType::Binary => pl.const_get::<_, Value>("Binary").unwrap(),
+            DataType::Array(inner, size) => {
+                let inner = Wrap(*inner);
+                let list_class = pl.const_get::<_, Value>("Array").unwrap();
+                list_class
+                    .funcall::<_, _, Value>("new", (size, inner))
+                    .unwrap()
+            }
             DataType::List(inner) => {
                 let inner = Wrap(*inner);
                 let list_class = pl.const_get::<_, Value>("List").unwrap();
@@ -321,11 +314,11 @@ impl IntoValue for Wrap<&StructChunked> {
 impl IntoValue for Wrap<&DurationChunked> {
     fn into_value_with(self, _: &RubyHandle) -> Value {
         let utils = utils();
-        let tu = Wrap(self.0.time_unit()).into_value();
+        let time_unit = Wrap(self.0.time_unit()).into_value();
         let iter = self.0.into_iter().map(|opt_v| {
             opt_v.map(|v| {
                 utils
-                    .funcall::<_, _, Value>("_to_ruby_duration", (v, tu))
+                    .funcall::<_, _, Value>("_to_ruby_duration", (v, time_unit))
                     .unwrap()
             })
         });
@@ -333,6 +326,64 @@ impl IntoValue for Wrap<&DurationChunked> {
     }
 }
 
+impl IntoValue for Wrap<&DatetimeChunked> {
+    fn into_value_with(self, _: &RubyHandle) -> Value {
+        let utils = utils();
+        let time_unit = Wrap(self.0.time_unit()).into_value();
+        let time_zone = self.0.time_zone().clone().into_value();
+        let iter = self.0.into_iter().map(|opt_v| {
+            opt_v.map(|v| {
+                utils
+                    .funcall::<_, _, Value>("_to_ruby_datetime", (v, time_unit, time_zone))
+                    .unwrap()
+            })
+        });
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&TimeChunked> {
+    fn into_value_with(self, _: &RubyHandle) -> Value {
+        let utils = utils();
+        let iter = self.0.into_iter().map(|opt_v| {
+            opt_v.map(|v| utils.funcall::<_, _, Value>("_to_ruby_time", (v,)).unwrap())
+        });
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&DateChunked> {
+    fn into_value_with(self, _: &RubyHandle) -> Value {
+        let utils = utils();
+        let iter = self.0.into_iter().map(|opt_v| {
+            opt_v.map(|v| utils.funcall::<_, _, Value>("_to_ruby_date", (v,)).unwrap())
+        });
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&DecimalChunked> {
+    fn into_value_with(self, _: &RubyHandle) -> Value {
+        let utils = utils();
+        let rb_scale = (-(self.0.scale() as i32)).into_value();
+        let iter = self.0.into_iter().map(|opt_v| {
+            opt_v.map(|v| {
+                utils
+                    .funcall::<_, _, Value>("_to_ruby_decimal", (v.to_string(), rb_scale))
+                    .unwrap()
+            })
+        });
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+fn abs_decimal_from_digits(digits: String, exp: i32) -> Option<(i128, usize)> {
+    match digits.parse::<i128>() {
+        Ok(v) => Some((v, ((digits.len() as i32) - exp) as usize)),
+        Err(_) => None,
+    }
+}
+
 impl TryConvert for Wrap<Field> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let name: String = ob.funcall("name", ())?;
@@ -362,10 +413,11 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::Datetime" => DataType::Datetime(TimeUnit::Microseconds, None),
                 "Polars::Time" => DataType::Time,
                 "Polars::Duration" => DataType::Duration(TimeUnit::Microseconds),
+                "Polars::Decimal" => DataType::Decimal(None, None),
                 "Polars::Float32" => DataType::Float32,
                 "Polars::Float64" => DataType::Float64,
-                // "Polars::Object" => DataType::Object(OBJECT_NAME),
-                "Polars::List" => DataType::List(Box::new(DataType::Boolean)),
+                "Polars::Object" => DataType::Object(OBJECT_NAME),
+                "Polars::List" => DataType::List(Box::new(DataType::Null)),
                 "Polars::Null" => DataType::Null,
                 "Polars::Unknown" => DataType::Unknown,
                 dt => {
@@ -378,6 +430,28 @@ impl TryConvert for Wrap<DataType> {
         } else if ob.try_convert::<String>().is_err() {
             let name = unsafe { ob.class().name() }.into_owned();
             match name.as_str() {
+                "Polars::Duration" => {
+                    let time_unit: Value = ob.funcall("time_unit", ()).unwrap();
+                    let time_unit = time_unit.try_convert::<Wrap<TimeUnit>>()?.0;
+                    DataType::Duration(time_unit)
+                }
+                "Polars::Datetime" => {
+                    let time_unit: Value = ob.funcall("time_unit", ()).unwrap();
+                    let time_unit = time_unit.try_convert::<Wrap<TimeUnit>>()?.0;
+                    let time_zone: Value = ob.funcall("time_zone", ()).unwrap();
+                    let time_zone = time_zone.try_convert()?;
+                    DataType::Datetime(time_unit, time_zone)
+                }
+                "Polars::Decimal" => {
+                    let precision = ob.funcall::<_, _, Value>("precision", ())?.try_convert()?;
+                    let scale = ob.funcall::<_, _, Value>("scale", ())?.try_convert()?;
+                    DataType::Decimal(precision, Some(scale))
+                }
+                "Polars::List" => {
+                    let inner: Value = ob.funcall("inner", ()).unwrap();
+                    let inner = inner.try_convert::<Wrap<DataType>>()?;
+                    DataType::List(Box::new(inner.0))
+                }
                 "Polars::Struct" => {
                     let arr: RArray = ob.funcall("fields", ())?;
                     let mut fields = Vec::with_capacity(arr.len());
@@ -386,7 +460,12 @@ impl TryConvert for Wrap<DataType> {
                     }
                     DataType::Struct(fields)
                 }
-                _ => todo!(),
+                dt => {
+                    return Err(RbTypeError::new_err(format!(
+                        "A {dt} object is not a correct polars DataType. \
+                        Hint: use the class without instantiating it.",
+                    )))
+                }
             }
         } else {
             match ob.try_convert::<String>()?.as_str() {
@@ -430,7 +509,7 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
             Ok(AnyValue::Boolean(ob.try_convert::<bool>()?).into())
         } else if let Some(v) = Integer::from_value(ob) {
             Ok(AnyValue::Int64(v.to_i64()?).into())
-        } else if let Some(v) = RFloat::from_value(ob) {
+        } else if let Some(v) = Float::from_value(ob) {
             Ok(AnyValue::Float64(v.to_f64()).into())
         } else if let Some(v) = RString::from_value(ob) {
             if v.enc_get() == Index::utf8() {
@@ -443,7 +522,8 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
             let sec = ob.funcall::<_, _, i64>("to_i", ())?;
             let nsec = ob.funcall::<_, _, i64>("nsec", ())?;
             let v = sec * 1_000_000_000 + nsec;
-            // TODO support time zone
+            // TODO support time zone when possible
+            // https://github.com/pola-rs/polars/issues/9103
             Ok(AnyValue::Datetime(v, TimeUnit::Nanoseconds, &None).into())
         } else if ob.is_nil() {
             Ok(AnyValue::Null.into())
@@ -464,13 +544,23 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
             if v.is_empty() {
                 Ok(Wrap(AnyValue::List(Series::new_empty("", &DataType::Null))))
             } else {
-                let avs = v.try_convert::<Wrap<Row>>()?.0 .0;
-                // use first `n` values to infer datatype
-                // this value is not too large as this will be done with every
-                // anyvalue that has to be converted, which can be many
-                let n = 25;
-                let dtype = any_values_to_dtype(&avs[..std::cmp::min(avs.len(), n)])
-                    .map_err(RbPolarsErr::from)?;
+                let list = v;
+
+                let mut avs = Vec::with_capacity(25);
+                let mut iter = list.each();
+
+                for item in (&mut iter).take(25) {
+                    avs.push(item?.try_convert::<Wrap<AnyValue>>()?.0)
+                }
+
+                let (dtype, _n_types) = any_values_to_dtype(&avs).map_err(RbPolarsErr::from)?;
+
+                // push the rest
+                avs.reserve(list.len());
+                for item in iter {
+                    avs.push(item?.try_convert::<Wrap<AnyValue>>()?.0)
+                }
+
                 let s = Series::from_any_values_and_dtype("", &avs, &dtype, true)
                     .map_err(RbPolarsErr::from)?;
                 Ok(Wrap(AnyValue::List(s)))
@@ -478,7 +568,11 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
         } else if ob.is_kind_of(crate::rb_modules::datetime()) {
             let sec: i64 = ob.funcall("to_i", ())?;
             let nsec: i64 = ob.funcall("nsec", ())?;
-            Ok(Wrap(AnyValue::Datetime(sec * 1_000_000_000 + nsec, TimeUnit::Nanoseconds, &None)))
+            Ok(Wrap(AnyValue::Datetime(
+                sec * 1_000_000_000 + nsec,
+                TimeUnit::Nanoseconds,
+                &None,
+            )))
         } else if ob.is_kind_of(crate::rb_modules::date()) {
             // convert to DateTime for UTC
             let v = ob
@@ -486,6 +580,20 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
                 .funcall::<_, _, Value>("to_time", ())?
                 .funcall::<_, _, i64>("to_i", ())?;
             Ok(Wrap(AnyValue::Date((v / 86400) as i32)))
+        } else if ob.is_kind_of(crate::rb_modules::bigdecimal()) {
+            let (sign, digits, _, exp): (i8, String, i32, i32) = ob
+                .funcall::<_, _, Value>("split", ())
+                .unwrap()
+                .try_convert()
+                .unwrap();
+            let (mut v, scale) = abs_decimal_from_digits(digits, exp).ok_or_else(|| {
+                RbPolarsErr::other("BigDecimal is too large to fit in Decimal128".into())
+            })?;
+            if sign < 0 {
+                // TODO better error
+                v = v.checked_neg().unwrap();
+            }
+            Ok(Wrap(AnyValue::Decimal(v, scale)))
         } else {
             Err(RbPolarsErr::other(format!(
                 "object type not supported {:?}",

data/ext/polars/src/dataframe.rs

@@ -456,12 +456,14 @@ impl RbDataFrame {
                 .finish(&mut self.df.borrow_mut())
                 .map_err(RbPolarsErr::from)?;
         } else {
-            let mut buf = get_file_like(rb_f, true)?;
-
+            let mut buf = Cursor::new(Vec::new());
             IpcWriter::new(&mut buf)
                 .with_compression(compression.0)
                 .finish(&mut self.df.borrow_mut())
                 .map_err(RbPolarsErr::from)?;
+            // TODO less copying
+            let rb_str = RString::from_slice(&buf.into_inner());
+            rb_f.funcall::<_, _, Value>("write", (rb_str,))?;
         }
         Ok(())
     }
@@ -794,22 +796,6 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
-    pub fn sort(&self, by_column: String, reverse: bool, nulls_last: bool) -> RbResult<Self> {
-        let df = self
-            .df
-            .borrow()
-            .sort_with_options(
-                &by_column,
-                SortOptions {
-                    descending: reverse,
-                    nulls_last,
-                    multithreaded: true,
-                },
-            )
-            .map_err(RbPolarsErr::from)?;
-        Ok(RbDataFrame::new(df))
-    }
-
     pub fn replace(&self, column: String, new_col: &RbSeries) -> RbResult<()> {
         self.df
             .borrow_mut()
@@ -933,11 +919,16 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
-    pub fn partition_by(&self, groups: Vec<String>, stable: bool) -> RbResult<RArray> {
-        let out = if stable {
-            self.df.borrow().partition_by_stable(groups)
+    pub fn partition_by(
+        &self,
+        by: Vec<String>,
+        maintain_order: bool,
+        include_key: bool,
+    ) -> RbResult<RArray> {
+        let out = if maintain_order {
+            self.df.borrow().partition_by_stable(by, include_key)
         } else {
-            self.df.borrow().partition_by(groups)
+            self.df.borrow().partition_by(by, include_key)
         }
         .map_err(RbPolarsErr::from)?;
         Ok(RArray::from_iter(out.into_iter().map(RbDataFrame::new)))
@@ -1024,13 +1015,18 @@ impl RbDataFrame {
         &self,
         columns: Option<Vec<String>>,
         separator: Option<String>,
+        drop_first: bool,
     ) -> RbResult<Self> {
         let df = match columns {
             Some(cols) => self.df.borrow().columns_to_dummies(
                 cols.iter().map(|x| x as &str).collect(),
                 separator.as_deref(),
+                drop_first,
             ),
-            None => self.df.borrow().to_dummies(separator.as_deref()),
+            None => self
+                .df
+                .borrow()
+                .to_dummies(separator.as_deref(), drop_first),
         }
         .map_err(RbPolarsErr::from)?;
         Ok(df.into())

data/ext/polars/src/error.rs

@@ -24,6 +24,14 @@ impl RbPolarsErr {
     }
 }
 
+pub struct RbTypeError {}
+
+impl RbTypeError {
+    pub fn new_err(message: String) -> Error {
+        Error::new(exception::type_error(), message)
+    }
+}
+
 pub struct RbValueError {}
 
 impl RbValueError {

data/ext/polars/src/expr/array.rs

@@ -0,0 +1,15 @@
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn array_max(&self) -> Self {
+        self.inner.clone().arr().max().into()
+    }
+
+    pub fn array_min(&self) -> Self {
+        self.inner.clone().arr().min().into()
+    }
+
+    pub fn array_sum(&self) -> Self {
+        self.inner.clone().arr().sum().into()
+    }
+}