polars-df 0.11.0 → 0.12.0

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.11.0"
+version = "0.12.0"
 license = "MIT"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
@@ -14,15 +14,15 @@ crate-type = ["cdylib"]
 ahash = "0.8"
 chrono = "0.4"
 either = "1.8"
-magnus = "0.6"
-polars-core = "=0.40.0"
-polars-parquet = "=0.40.0"
-polars-utils = "=0.40.0"
+magnus = "0.7"
+polars-core = "=0.41.3"
+polars-parquet = "=0.41.3"
+polars-utils = "=0.41.3"
 serde_json = "1"
 smartstring = "1"
 
 [dependencies.polars]
-version = "=0.40.0"
+version = "=0.41.3"
 features = [
     "abs",
     "approx_unique",
@@ -41,7 +41,6 @@ features = [
     "cumulative_eval",
     "cutqcut",
     "dataframe_arithmetic",
-    "date_offset",
     "diagonal_concat",
     "diff",
     "dot_product",
@@ -56,6 +55,7 @@ features = [
     "interpolate",
     "ipc",
     "ipc_streaming",
+    "is_between",
     "is_first_distinct",
     "is_in",
     "is_last_distinct",
@@ -74,6 +74,9 @@ features = [
     "meta",
     "mode",
     "moment",
+    "month_start",
+    "month_end",
+    "offset_by",
     "object",
     "parquet",
     "partition_by",

data/ext/polars/src/batched_csv.rs

@@ -95,7 +95,7 @@ impl RbBatchedCsv {
             .with_projection(projection.map(Arc::new))
             .with_rechunk(rechunk)
             .with_chunk_size(chunk_size)
-            .with_columns(columns.map(Arc::new))
+            .with_columns(columns.map(Arc::from))
             .with_n_threads(n_threads)
             .with_dtype_overwrite(overwrite_dtype_slice.map(Arc::new))
             .with_low_memory(low_memory)

data/ext/polars/src/conversion/any_value.rs

@@ -0,0 +1,261 @@
+use magnus::encoding::{EncodingCapable, Index};
+use magnus::{
+    class, prelude::*, r_hash::ForEach, IntoValue, RArray, RHash, RString, Ruby, TryConvert, Value,
+};
+use polars::prelude::*;
+use polars_core::utils::any_values_to_supertype_and_n_dtypes;
+
+use super::{struct_dict, ObjectValue, Wrap};
+
+use crate::error::RbOverflowError;
+use crate::rb_modules::utils;
+use crate::{RbPolarsErr, RbResult, RbSeries};
+
+impl IntoValue for Wrap<AnyValue<'_>> {
+    fn into_value_with(self, ruby: &Ruby) -> Value {
+        any_value_into_rb_object(self.0, ruby)
+    }
+}
+
+impl<'s> TryConvert for Wrap<AnyValue<'s>> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        rb_object_to_any_value(ob, true).map(Wrap)
+    }
+}
+
+pub(crate) fn any_value_into_rb_object(av: AnyValue, ruby: &Ruby) -> Value {
+    match av {
+        AnyValue::UInt8(v) => ruby.into_value(v),
+        AnyValue::UInt16(v) => ruby.into_value(v),
+        AnyValue::UInt32(v) => ruby.into_value(v),
+        AnyValue::UInt64(v) => ruby.into_value(v),
+        AnyValue::Int8(v) => ruby.into_value(v),
+        AnyValue::Int16(v) => ruby.into_value(v),
+        AnyValue::Int32(v) => ruby.into_value(v),
+        AnyValue::Int64(v) => ruby.into_value(v),
+        AnyValue::Float32(v) => ruby.into_value(v),
+        AnyValue::Float64(v) => ruby.into_value(v),
+        AnyValue::Null => ruby.qnil().as_value(),
+        AnyValue::Boolean(v) => ruby.into_value(v),
+        AnyValue::String(v) => ruby.into_value(v),
+        AnyValue::StringOwned(v) => ruby.into_value(v.as_str()),
+        AnyValue::Categorical(idx, rev, arr) | AnyValue::Enum(idx, rev, arr) => {
+            let s = if arr.is_null() {
+                rev.get(idx)
+            } else {
+                unsafe { arr.deref_unchecked().value(idx as usize) }
+            };
+            s.into_value()
+        }
+        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::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) => {
+            let object = v.as_any().downcast_ref::<ObjectValue>().unwrap();
+            object.to_object()
+        }
+        AnyValue::ObjectOwned(v) => {
+            let object = v.0.as_any().downcast_ref::<ObjectValue>().unwrap();
+            object.to_object()
+        }
+        AnyValue::Binary(v) => RString::from_slice(v).into_value(),
+        AnyValue::BinaryOwned(v) => RString::from_slice(&v).into_value(),
+        AnyValue::Decimal(v, scale) => utils()
+            .funcall("_to_ruby_decimal", (v.to_string(), -(scale as i32)))
+            .unwrap(),
+    }
+}
+
+pub(crate) fn rb_object_to_any_value<'s>(ob: Value, strict: bool) -> RbResult<AnyValue<'s>> {
+    // Conversion functions.
+    fn get_null(_ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        Ok(AnyValue::Null)
+    }
+
+    fn get_bool(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        let b = bool::try_convert(ob)?;
+        Ok(AnyValue::Boolean(b))
+    }
+
+    fn get_int(ob: Value, strict: bool) -> RbResult<AnyValue<'static>> {
+        if let Ok(v) = i64::try_convert(ob) {
+            Ok(AnyValue::Int64(v))
+        } else if let Ok(v) = u64::try_convert(ob) {
+            Ok(AnyValue::UInt64(v))
+        } else if !strict {
+            let f = f64::try_convert(ob)?;
+            Ok(AnyValue::Float64(f))
+        } else {
+            Err(RbOverflowError::new_err(format!(
+                "int value too large for Polars integer types: {ob}"
+            )))
+        }
+    }
+
+    fn get_float(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        Ok(AnyValue::Float64(f64::try_convert(ob)?))
+    }
+
+    fn get_str(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        let v = RString::from_value(ob).unwrap();
+        if v.enc_get() == Index::utf8() {
+            Ok(AnyValue::StringOwned(v.to_string()?.into()))
+        } else {
+            Ok(AnyValue::BinaryOwned(unsafe { v.as_slice() }.to_vec()))
+        }
+    }
+
+    fn get_list(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        let v = RArray::from_value(ob).unwrap();
+        if v.is_empty() {
+            Ok(AnyValue::List(Series::new_empty("", &DataType::Null)))
+        } else {
+            let list = v;
+
+            let mut avs = Vec::with_capacity(25);
+            let mut iter = list.into_iter();
+
+            for item in (&mut iter).take(25) {
+                avs.push(Wrap::<AnyValue>::try_convert(item)?.0)
+            }
+
+            let (dtype, _n_types) =
+                any_values_to_supertype_and_n_dtypes(&avs).map_err(RbPolarsErr::from)?;
+
+            // push the rest
+            avs.reserve(list.len());
+            for item in iter {
+                avs.push(Wrap::<AnyValue>::try_convert(item)?.0)
+            }
+
+            let s = Series::from_any_values_and_dtype("", &avs, &dtype, true)
+                .map_err(RbPolarsErr::from)?;
+            Ok(AnyValue::List(s))
+        }
+    }
+
+    fn get_list_from_series(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        let s = super::get_series(ob)?;
+        Ok(AnyValue::List(s))
+    }
+
+    fn get_struct(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        let dict = RHash::from_value(ob).unwrap();
+        let len = dict.len();
+        let mut keys = Vec::with_capacity(len);
+        let mut vals = Vec::with_capacity(len);
+        dict.foreach(|k: Value, v: Value| {
+            let key = String::try_convert(k)?;
+            let val = Wrap::<AnyValue>::try_convert(v)?.0;
+            let dtype = DataType::from(&val);
+            keys.push(Field::new(&key, dtype));
+            vals.push(val);
+            Ok(ForEach::Continue)
+        })?;
+        Ok(AnyValue::StructOwned(Box::new((vals, keys))))
+    }
+
+    fn get_date(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        // convert to DateTime for UTC
+        let v = ob
+            .funcall::<_, _, Value>("to_datetime", ())?
+            .funcall::<_, _, Value>("to_time", ())?
+            .funcall::<_, _, i64>("to_i", ())?;
+        Ok(AnyValue::Date((v / 86400) as i32))
+    }
+
+    fn get_time(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        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 when possible
+        // https://github.com/pola-rs/polars/issues/9103
+        Ok(AnyValue::Datetime(v, TimeUnit::Nanoseconds, &None))
+    }
+
+    fn get_datetime(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        let sec: i64 = ob.funcall("to_i", ())?;
+        let nsec: i64 = ob.funcall("nsec", ())?;
+        Ok(AnyValue::Datetime(
+            sec * 1_000_000_000 + nsec,
+            TimeUnit::Nanoseconds,
+            &None,
+        ))
+    }
+
+    fn get_decimal(ob: Value, _strict: bool) -> RbResult<AnyValue<'static>> {
+        fn abs_decimal_from_digits(digits: String, exp: i32) -> Option<(i128, usize)> {
+            let exp = exp - (digits.len() as i32);
+            match digits.parse::<i128>() {
+                Ok(mut v) => {
+                    let scale = if exp > 0 {
+                        v = 10_i128
+                            .checked_pow(exp as u32)
+                            .and_then(|factor| v.checked_mul(factor))?;
+                        0
+                    } else {
+                        (-exp) as usize
+                    };
+                    Some((v, scale))
+                }
+                Err(_) => None,
+            }
+        }
+
+        let (sign, digits, _, exp): (i8, String, i32, i32) = ob.funcall("split", ()).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(AnyValue::Decimal(v, scale))
+    }
+
+    if ob.is_nil() {
+        get_null(ob, strict)
+    } else if ob.is_kind_of(class::true_class()) || ob.is_kind_of(class::false_class()) {
+        get_bool(ob, strict)
+    } else if ob.is_kind_of(class::integer()) {
+        get_int(ob, strict)
+    } else if ob.is_kind_of(class::float()) {
+        get_float(ob, strict)
+    } else if ob.is_kind_of(class::string()) {
+        get_str(ob, strict)
+    } else if ob.is_kind_of(class::array()) {
+        get_list(ob, strict)
+    } else if ob.is_kind_of(class::hash()) {
+        get_struct(ob, strict)
+    } else if ob.respond_to("_s", true)? {
+        get_list_from_series(ob, strict)
+    // call is_a? for ActiveSupport::TimeWithZone
+    } else if ob.funcall::<_, _, bool>("is_a?", (class::time(),))? {
+        get_time(ob, strict)
+    } else if ob.is_kind_of(crate::rb_modules::datetime()) {
+        get_datetime(ob, strict)
+    } else if ob.is_kind_of(crate::rb_modules::date()) {
+        get_date(ob, strict)
+    } else if ob.is_kind_of(crate::rb_modules::bigdecimal()) {
+        get_decimal(ob, strict)
+    } else {
+        Err(RbPolarsErr::other(format!(
+            "object type not supported {:?}",
+            ob
+        )))
+    }
+}

data/ext/polars/src/conversion/chunked_array.rs

@@ -11,8 +11,8 @@ impl TryConvert for Wrap<StringChunked> {
         let (seq, len) = get_rbseq(obj)?;
         let mut builder = StringChunkedBuilder::new("", len);
 
-        for res in seq.each() {
-            let item = res?;
+        for res in seq.into_iter() {
+            let item = res;
             match String::try_convert(item) {
                 Ok(val) => builder.append_value(&val),
                 Err(_) => builder.append_null(),
@@ -27,8 +27,8 @@ impl TryConvert for Wrap<BinaryChunked> {
         let (seq, len) = get_rbseq(obj)?;
         let mut builder = BinaryChunkedBuilder::new("", len);
 
-        for res in seq.each() {
-            let item = res?;
+        for res in seq.into_iter() {
+            let item = res;
             match RString::try_convert(item) {
                 Ok(val) => builder.append_value(unsafe { val.as_slice() }),
                 Err(_) => builder.append_null(),

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

@@ -1,4 +1,4 @@
-pub(crate) mod anyvalue;
+pub(crate) mod any_value;
 mod chunked_array;
 
 use std::fmt::{Debug, Display, Formatter};
@@ -7,7 +7,7 @@ use std::num::NonZeroUsize;
 
 use magnus::{
     class, exception, prelude::*, r_hash::ForEach, value::Opaque, IntoValue, Module, RArray, RHash,
-    Ruby, TryConvert, Value,
+    Ruby, Symbol, TryConvert, Value,
 };
 use polars::chunked_array::object::PolarsObjectSafe;
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
@@ -75,7 +75,7 @@ pub(crate) fn get_df(obj: Value) -> RbResult<DataFrame> {
 
 pub(crate) fn get_lf(obj: Value) -> RbResult<LazyFrame> {
     let rbdf = obj.funcall::<_, _, &RbLazyFrame>("_ldf", ())?;
-    Ok(rbdf.ldf.clone())
+    Ok(rbdf.ldf.borrow().clone())
 }
 
 pub(crate) fn get_series(obj: Value) -> RbResult<Series> {
@@ -379,8 +379,8 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::Struct" => {
                     let arr: RArray = ob.funcall("fields", ())?;
                     let mut fields = Vec::with_capacity(arr.len());
-                    for v in arr.each() {
-                        fields.push(Wrap::<Field>::try_convert(v?)?.0);
+                    for v in arr.into_iter() {
+                        fields.push(Wrap::<Field>::try_convert(v)?.0);
                     }
                     DataType::Struct(fields)
                 }
@@ -427,11 +427,36 @@ impl TryConvert for Wrap<DataType> {
     }
 }
 
+impl TryConvert for Wrap<StatisticsOptions> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let mut statistics = StatisticsOptions::empty();
+
+        let dict = RHash::try_convert(ob)?;
+        dict.foreach(|key: Symbol, val: bool| {
+            match key.name()?.as_ref() {
+                "min" => statistics.min_value = val,
+                "max" => statistics.max_value = val,
+                "distinct_count" => statistics.distinct_count = val,
+                "null_count" => statistics.null_count = val,
+                _ => {
+                    return Err(RbTypeError::new_err(format!(
+                        "'{key}' is not a valid statistic option",
+                    )))
+                }
+            }
+            Ok(ForEach::Continue)
+        })
+        .unwrap();
+
+        Ok(Wrap(statistics))
+    }
+}
+
 impl<'s> TryConvert for Wrap<Row<'s>> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let mut vals: Vec<Wrap<AnyValue<'s>>> = Vec::new();
-        for item in RArray::try_convert(ob)?.each() {
-            vals.push(Wrap::<AnyValue<'s>>::try_convert(item?)?);
+        for item in RArray::try_convert(ob)?.into_iter() {
+            vals.push(Wrap::<AnyValue<'s>>::try_convert(item)?);
         }
         let vals: Vec<AnyValue> = unsafe { std::mem::transmute(vals) };
         Ok(Wrap(Row(vals)))
@@ -686,14 +711,13 @@ impl TryConvert for Wrap<JoinType> {
         let parsed = match String::try_convert(ob)?.as_str() {
             "inner" => JoinType::Inner,
             "left" => JoinType::Left,
-            "outer" => JoinType::Outer,
-            "outer_coalesce" => JoinType::Outer,
+            "full" => JoinType::Full,
             "semi" => JoinType::Semi,
             "anti" => JoinType::Anti,
             "cross" => JoinType::Cross,
             v => {
                 return Err(RbValueError::new_err(format!(
-                "how must be one of {{'inner', 'left', 'outer', 'semi', 'anti', 'cross'}}, got {}",
+                "how must be one of {{'inner', 'left', 'full', 'semi', 'anti', 'cross'}}, got {}",
                 v
             )))
             }
@@ -903,6 +927,23 @@ impl TryConvert for Wrap<SearchSortedSide> {
     }
 }
 
+impl TryConvert for Wrap<ClosedInterval> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match String::try_convert(ob)?.as_str() {
+            "both" => ClosedInterval::Both,
+            "left" => ClosedInterval::Left,
+            "right" => ClosedInterval::Right,
+            "none" => ClosedInterval::None,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "`closed` must be one of {{'both', 'left', 'right', 'none'}}, got {v}",
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<WindowMapping> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match String::try_convert(ob)?.as_str() {

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

@@ -13,11 +13,11 @@ impl RbDataFrame {
         schema: Option<Wrap<Schema>>,
     ) -> RbResult<Self> {
         let mut data = Vec::with_capacity(rb_rows.len());
-        for v in rb_rows.each() {
-            let rb_row = RArray::try_convert(v?)?;
+        for v in rb_rows.into_iter() {
+            let rb_row = RArray::try_convert(v)?;
             let mut row = Vec::with_capacity(rb_row.len());
-            for val in rb_row.each() {
-                row.push(Wrap::<AnyValue>::try_convert(val?)?.0);
+            for val in rb_row.into_iter() {
+                row.push(Wrap::<AnyValue>::try_convert(val)?.0);
             }
             data.push(Row(row));
         }
@@ -130,8 +130,7 @@ where
 fn dicts_to_rows<'a>(data: &Value, names: &'a [String], _strict: bool) -> RbResult<Vec<Row<'a>>> {
     let (data, len) = get_rbseq(*data)?;
     let mut rows = Vec::with_capacity(len);
-    for d in data.each() {
-        let d = d?;
+    for d in data.into_iter() {
         let d = RHash::try_convert(d)?;
 
         let mut row = Vec::with_capacity(names.len());
@@ -170,8 +169,7 @@ fn infer_schema_names_from_data(
         .unwrap_or(data_len);
 
     let mut names = PlIndexSet::new();
-    for d in data.each().take(infer_schema_length) {
-        let d = d?;
+    for d in data.into_iter().take(infer_schema_length) {
         let d = RHash::try_convert(d)?;
         d.foreach(|name: Value, _value: Value| {
             if let Some(v) = Symbol::from_value(name) {

data/ext/polars/src/dataframe/general.rs

@@ -15,8 +15,8 @@ use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
 impl RbDataFrame {
     pub fn init(columns: RArray) -> RbResult<Self> {
         let mut cols = Vec::new();
-        for i in columns.each() {
-            cols.push(<&RbSeries>::try_convert(i?)?.series.borrow().clone());
+        for i in columns.into_iter() {
+            cols.push(<&RbSeries>::try_convert(i)?.series.borrow().clone());
         }
         let df = DataFrame::new(cols).map_err(RbPolarsErr::from)?;
         Ok(RbDataFrame::new(df))
@@ -341,45 +341,42 @@ impl RbDataFrame {
         RbDataFrame::new(self.df.borrow().clone())
     }
 
-    pub fn melt(
+    pub fn unpivot(
         &self,
-        id_vars: Vec<String>,
-        value_vars: Vec<String>,
+        on: Vec<String>,
+        index: Vec<String>,
         value_name: Option<String>,
         variable_name: Option<String>,
     ) -> RbResult<Self> {
-        let args = MeltArgs {
-            id_vars: strings_to_smartstrings(id_vars),
-            value_vars: strings_to_smartstrings(value_vars),
+        let args = UnpivotArgs {
+            on: strings_to_smartstrings(on),
+            index: strings_to_smartstrings(index),
             value_name: value_name.map(|s| s.into()),
             variable_name: variable_name.map(|s| s.into()),
             streamable: false,
         };
 
-        let df = self.df.borrow().melt2(args).map_err(RbPolarsErr::from)?;
+        let df = self.df.borrow().unpivot2(args).map_err(RbPolarsErr::from)?;
         Ok(RbDataFrame::new(df))
     }
 
     #[allow(clippy::too_many_arguments)]
     pub fn pivot_expr(
         &self,
-        index: Vec<String>,
-        columns: Vec<String>,
+        on: Vec<String>,
+        index: Option<Vec<String>>,
         values: Option<Vec<String>>,
         maintain_order: bool,
         sort_columns: bool,
         aggregate_expr: Option<&RbExpr>,
         separator: Option<String>,
     ) -> RbResult<Self> {
-        let fun = match maintain_order {
-            true => pivot_stable,
-            false => pivot,
-        };
+        let fun = if maintain_order { pivot_stable } else { pivot };
         let agg_expr = aggregate_expr.map(|aggregate_expr| aggregate_expr.inner.clone());
         let df = fun(
             &self.df.borrow(),
+            on,
             index,
-            columns,
             values,
             sort_columns,
             agg_expr,
@@ -569,23 +566,16 @@ impl RbDataFrame {
         by: Vec<String>,
         index_column: String,
         every: String,
-        offset: String,
         stable: bool,
     ) -> RbResult<Self> {
         let out = if stable {
-            self.df.borrow().upsample_stable(
-                by,
-                &index_column,
-                Duration::parse(&every),
-                Duration::parse(&offset),
-            )
+            self.df
+                .borrow()
+                .upsample_stable(by, &index_column, Duration::parse(&every))
         } else {
-            self.df.borrow().upsample(
-                by,
-                &index_column,
-                Duration::parse(&every),
-                Duration::parse(&offset),
-            )
+            self.df
+                .borrow()
+                .upsample(by, &index_column, Duration::parse(&every))
         };
         let out = out.map_err(RbPolarsErr::from)?;
         Ok(out.into())