polars-df 0.8.0 → 0.10.0

data/ext/polars/Cargo.toml

@@ -1,10 +1,10 @@
 [package]
 name = "polars"
-version = "0.8.0"
+version = "0.10.0"
 license = "MIT"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
-rust-version = "1.73.0"
+rust-version = "1.76.0"
 publish = false
 
 [lib]
@@ -15,23 +15,26 @@ ahash = "0.8"
 chrono = "0.4"
 either = "1.8"
 magnus = "0.6"
-polars-core = "=0.36.2"
-polars-parquet = "=0.36.2"
-polars-utils = "=0.36.2"
+polars-core = "=0.39.2"
+polars-parquet = "=0.39.2"
+polars-utils = "=0.39.2"
 serde_json = "1"
 smartstring = "1"
 
 [dependencies.polars]
-version = "=0.36.2"
+version = "=0.39.2"
 features = [
     "abs",
     "approx_unique",
     "arg_where",
+    "array_any_all",
+    "array_count",
     "asof_join",
     "avro",
     "binary_encoding",
     "concat_str",
     "cov",
+    "cross_join",
     "cse",
     "csv",
     "cum_agg",
@@ -45,9 +48,11 @@ features = [
     "dtype-full",
     "dynamic_group_by",
     "ewma",
+    "ewma_by",
+    "extract_groups",
     "extract_jsonpath",
+    "find_many",
     "fmt",
-    "horizontal_concat",
     "interpolate",
     "ipc",
     "is_first_distinct",
@@ -56,12 +61,15 @@ features = [
     "is_unique",
     "json",
     "lazy",
-    "lazy_regex",
+    "list_any_all",
     "list_count",
+    "list_drop_nulls",
     "list_eval",
     "list_gather",
+    "list_sample",
     "list_to_struct",
     "log",
+    "merge_sorted",
     "meta",
     "mode",
     "moment",
@@ -77,8 +85,10 @@ features = [
     "random",
     "rank",
     "range",
+    "regex",
     "reinterpret",
     "repeat_by",
+    "replace",
     "rle",
     "rolling_window",
     "round_series",

data/ext/polars/src/batched_csv.rs

@@ -1,6 +1,6 @@
 use magnus::{prelude::*, RArray, Value};
 use polars::io::mmap::MmapBytesReader;
-use polars::io::RowCount;
+use polars::io::RowIndex;
 use polars::prelude::read_impl::OwnedBatchedCsvReader;
 use polars::prelude::*;
 use std::cell::RefCell;
@@ -46,15 +46,16 @@ impl RbBatchedCsv {
         let null_values = Option::<Wrap<NullValues>>::try_convert(arguments[18])?;
         let try_parse_dates = bool::try_convert(arguments[19])?;
         let skip_rows_after_header = usize::try_convert(arguments[20])?;
-        let row_count = Option::<(String, IdxSize)>::try_convert(arguments[21])?;
+        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[21])?;
         let sample_size = usize::try_convert(arguments[22])?;
         let eol_char = String::try_convert(arguments[23])?;
+        let truncate_ragged_lines = bool::try_convert(arguments[24])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
         let eol_char = eol_char.as_bytes()[0];
 
-        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
 
         let quote_char = if let Some(s) = quote_char {
             if s.is_empty() {
@@ -106,8 +107,9 @@ impl RbBatchedCsv {
             .with_quote_char(quote_char)
             .with_end_of_line_char(eol_char)
             .with_skip_rows_after_header(skip_rows_after_header)
-            .with_row_count(row_count)
-            .sample_size(sample_size);
+            .with_row_index(row_index)
+            .sample_size(sample_size)
+            .truncate_ragged_lines(truncate_ragged_lines);
 
         let reader = if low_memory {
             let reader = reader

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

@@ -0,0 +1,186 @@
+use magnus::encoding::{EncodingCapable, Index};
+use magnus::{
+    class, prelude::*, r_hash::ForEach, Float, Integer, 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::rb_modules::utils;
+use crate::{RbPolarsErr, RbResult, RbSeries};
+
+impl IntoValue for Wrap<AnyValue<'_>> {
+    fn into_value_with(self, ruby: &Ruby) -> Value {
+        match self.0 {
+            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(),
+        }
+    }
+}
+
+impl<'s> TryConvert for Wrap<AnyValue<'s>> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        if ob.is_kind_of(class::true_class()) || ob.is_kind_of(class::false_class()) {
+            Ok(AnyValue::Boolean(bool::try_convert(ob)?).into())
+        } else if let Some(v) = Integer::from_value(ob) {
+            Ok(AnyValue::Int64(v.to_i64()?).into())
+        } 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() {
+                Ok(AnyValue::StringOwned(v.to_string()?.into()).into())
+            } else {
+                Ok(AnyValue::BinaryOwned(unsafe { v.as_slice() }.to_vec()).into())
+            }
+        // call is_a? for ActiveSupport::TimeWithZone
+        } else if ob.funcall::<_, _, bool>("is_a?", (class::time(),))? {
+            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).into())
+        } else if ob.is_nil() {
+            Ok(AnyValue::Null.into())
+        } else if let Some(dict) = RHash::from_value(ob) {
+            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(Wrap(AnyValue::StructOwned(Box::new((vals, keys)))))
+        } else if let Some(v) = RArray::from_value(ob) {
+            if v.is_empty() {
+                Ok(Wrap(AnyValue::List(Series::new_empty("", &DataType::Null))))
+            } else {
+                let list = v;
+
+                let mut avs = Vec::with_capacity(25);
+                let mut iter = list.each();
+
+                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(Wrap(AnyValue::List(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,
+            )))
+        } else if ob.is_kind_of(crate::rb_modules::date()) {
+            // convert to DateTime for UTC
+            let v = ob
+                .funcall::<_, _, Value>("to_datetime", ())?
+                .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("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(Wrap(AnyValue::Decimal(v, scale)))
+        } else {
+            Err(RbPolarsErr::other(format!(
+                "object type not supported {:?}",
+                ob
+            )))
+        }
+    }
+}
+
+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,
+    }
+}

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

@@ -0,0 +1,140 @@
+use magnus::{prelude::*, IntoValue, RArray, RString, Ruby, TryConvert, Value};
+use polars::prelude::*;
+
+use super::{get_rbseq, struct_dict, Wrap};
+
+use crate::rb_modules::utils;
+use crate::RbResult;
+
+impl TryConvert for Wrap<StringChunked> {
+    fn try_convert(obj: Value) -> RbResult<Self> {
+        let (seq, len) = get_rbseq(obj)?;
+        let mut builder = StringChunkedBuilder::new("", len);
+
+        for res in seq.each() {
+            let item = res?;
+            match String::try_convert(item) {
+                Ok(val) => builder.append_value(&val),
+                Err(_) => builder.append_null(),
+            }
+        }
+        Ok(Wrap(builder.finish()))
+    }
+}
+
+impl TryConvert for Wrap<BinaryChunked> {
+    fn try_convert(obj: Value) -> RbResult<Self> {
+        let (seq, len) = get_rbseq(obj)?;
+        let mut builder = BinaryChunkedBuilder::new("", len);
+
+        for res in seq.each() {
+            let item = res?;
+            match RString::try_convert(item) {
+                Ok(val) => builder.append_value(unsafe { val.as_slice() }),
+                Err(_) => builder.append_null(),
+            }
+        }
+        Ok(Wrap(builder.finish()))
+    }
+}
+
+impl IntoValue for Wrap<&StringChunked> {
+    fn into_value_with(self, _: &Ruby) -> Value {
+        let iter = self.0.into_iter();
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&BinaryChunked> {
+    fn into_value_with(self, _: &Ruby) -> Value {
+        let iter = self
+            .0
+            .into_iter()
+            .map(|opt_bytes| opt_bytes.map(RString::from_slice));
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&StructChunked> {
+    fn into_value_with(self, _: &Ruby) -> Value {
+        let s = self.0.clone().into_series();
+        // todo! iterate its chunks and flatten.
+        // make series::iter() accept a chunk index.
+        let s = s.rechunk();
+        let iter = s.iter().map(|av| {
+            if let AnyValue::Struct(_, _, flds) = av {
+                struct_dict(av._iter_struct_av(), flds)
+            } else {
+                unreachable!()
+            }
+        });
+
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&DurationChunked> {
+    fn into_value_with(self, _: &Ruby) -> Value {
+        let utils = utils();
+        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, time_unit))
+                    .unwrap()
+            })
+        });
+        RArray::from_iter(iter).into_value()
+    }
+}
+
+impl IntoValue for Wrap<&DatetimeChunked> {
+    fn into_value_with(self, _: &Ruby) -> 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, _: &Ruby) -> 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, _: &Ruby) -> 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, _: &Ruby) -> 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()
+    }
+}