polars-df 0.10.0 → 0.12.0

data/README.md

@@ -1,8 +1,8 @@
-# Polars Ruby
+# Ruby Polars
 
 :fire: Blazingly fast DataFrames for Ruby, powered by [Polars](https://github.com/pola-rs/polars)
 
-[![Build Status](https://github.com/ankane/polars-ruby/actions/workflows/build.yml/badge.svg)](https://github.com/ankane/polars-ruby/actions)
+[![Build Status](https://github.com/ankane/ruby-polars/actions/workflows/build.yml/badge.svg)](https://github.com/ankane/ruby-polars/actions)
 
 ## Installation
 
@@ -420,16 +420,16 @@ View the [changelog](CHANGELOG.md)
 
 Everyone is encouraged to help improve this project. Here are a few ways you can help:
 
-- [Report bugs](https://github.com/ankane/polars-ruby/issues)
-- Fix bugs and [submit pull requests](https://github.com/ankane/polars-ruby/pulls)
+- [Report bugs](https://github.com/ankane/ruby-polars/issues)
+- Fix bugs and [submit pull requests](https://github.com/ankane/ruby-polars/pulls)
 - Write, clarify, or fix documentation
 - Suggest or add new features
 
 To get started with development:
 
 ```sh
-git clone https://github.com/ankane/polars-ruby.git
-cd polars-ruby
+git clone https://github.com/ankane/ruby-polars.git
+cd ruby-polars
 bundle install
 bundle exec rake compile
 bundle exec rake test

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.10.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.39.2"
-polars-parquet = "=0.39.2"
-polars-utils = "=0.39.2"
+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.39.2"
+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",
@@ -55,6 +54,8 @@ features = [
     "fmt",
     "interpolate",
     "ipc",
+    "ipc_streaming",
+    "is_between",
     "is_first_distinct",
     "is_in",
     "is_last_distinct",
@@ -73,6 +74,9 @@ features = [
     "meta",
     "mode",
     "moment",
+    "month_start",
+    "month_end",
+    "offset_by",
     "object",
     "parquet",
     "partition_by",
@@ -91,6 +95,7 @@ features = [
     "replace",
     "rle",
     "rolling_window",
+    "rolling_window_by",
     "round_series",
     "row_hash",
     "search_sorted",

data/ext/polars/src/batched_csv.rs

@@ -1,23 +1,19 @@
+use std::cell::RefCell;
+use std::path::PathBuf;
+use std::sync::Mutex;
+
 use magnus::{prelude::*, RArray, Value};
+use polars::io::csv::read::OwnedBatchedCsvReader;
 use polars::io::mmap::MmapBytesReader;
 use polars::io::RowIndex;
-use polars::prelude::read_impl::OwnedBatchedCsvReader;
 use polars::prelude::*;
-use std::cell::RefCell;
-use std::path::PathBuf;
 
 use crate::conversion::*;
-use crate::prelude::read_impl::OwnedBatchedCsvReaderMmap;
 use crate::{RbDataFrame, RbPolarsErr, RbResult};
 
-pub enum BatchedReader {
-    MMap(OwnedBatchedCsvReaderMmap),
-    Read(OwnedBatchedCsvReader),
-}
-
 #[magnus::wrap(class = "Polars::RbBatchedCsv")]
 pub struct RbBatchedCsv {
-    pub reader: RefCell<BatchedReader>,
+    pub reader: RefCell<Mutex<OwnedBatchedCsvReader>>,
 }
 
 impl RbBatchedCsv {
@@ -44,19 +40,23 @@ impl RbBatchedCsv {
         let comment_prefix = Option::<String>::try_convert(arguments[16])?;
         let quote_char = Option::<String>::try_convert(arguments[17])?;
         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_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])?;
+        let missing_utf8_is_empty_string = bool::try_convert(arguments[19])?;
+        let try_parse_dates = bool::try_convert(arguments[20])?;
+        let skip_rows_after_header = usize::try_convert(arguments[21])?;
+        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[22])?;
+        let sample_size = usize::try_convert(arguments[23])?;
+        let eol_char = String::try_convert(arguments[24])?;
+        let raise_if_empty = bool::try_convert(arguments[25])?;
+        let truncate_ragged_lines = bool::try_convert(arguments[26])?;
+        let decimal_comma = bool::try_convert(arguments[27])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
         let eol_char = eol_char.as_bytes()[0];
-
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
-
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
         let quote_char = if let Some(s) = quote_char {
             if s.is_empty() {
                 None
@@ -86,54 +86,55 @@ impl RbBatchedCsv {
 
         let file = std::fs::File::open(path).map_err(RbPolarsErr::io)?;
         let reader = Box::new(file) as Box<dyn MmapBytesReader>;
-        let reader = CsvReader::new(reader)
-            .infer_schema(infer_schema_length)
-            .has_header(has_header)
+        let reader = CsvReadOptions::default()
+            .with_infer_schema_length(infer_schema_length)
+            .with_has_header(has_header)
             .with_n_rows(n_rows)
-            .with_separator(separator.as_bytes()[0])
             .with_skip_rows(skip_rows)
             .with_ignore_errors(ignore_errors)
-            .with_projection(projection)
+            .with_projection(projection.map(Arc::new))
             .with_rechunk(rechunk)
             .with_chunk_size(chunk_size)
-            .with_encoding(encoding.0)
-            .with_columns(columns)
+            .with_columns(columns.map(Arc::from))
             .with_n_threads(n_threads)
-            .with_dtypes_slice(overwrite_dtype_slice.as_deref())
-            .low_memory(low_memory)
-            .with_comment_prefix(comment_prefix.as_deref())
-            .with_null_values(null_values)
-            .with_try_parse_dates(try_parse_dates)
-            .with_quote_char(quote_char)
-            .with_end_of_line_char(eol_char)
+            .with_dtype_overwrite(overwrite_dtype_slice.map(Arc::new))
+            .with_low_memory(low_memory)
             .with_skip_rows_after_header(skip_rows_after_header)
             .with_row_index(row_index)
-            .sample_size(sample_size)
-            .truncate_ragged_lines(truncate_ragged_lines);
+            .with_sample_size(sample_size)
+            .with_raise_if_empty(raise_if_empty)
+            .with_parse_options(
+                CsvParseOptions::default()
+                    .with_separator(separator.as_bytes()[0])
+                    .with_encoding(encoding.0)
+                    .with_missing_is_null(!missing_utf8_is_empty_string)
+                    .with_comment_prefix(comment_prefix.as_deref())
+                    .with_null_values(null_values)
+                    .with_try_parse_dates(try_parse_dates)
+                    .with_quote_char(quote_char)
+                    .with_eol_char(eol_char)
+                    .with_truncate_ragged_lines(truncate_ragged_lines)
+                    .with_decimal_comma(decimal_comma),
+            )
+            .into_reader_with_file_handle(reader);
 
-        let reader = if low_memory {
-            let reader = reader
-                .batched_read(overwrite_dtype.map(Arc::new))
-                .map_err(RbPolarsErr::from)?;
-            BatchedReader::Read(reader)
-        } else {
-            let reader = reader
-                .batched_mmap(overwrite_dtype.map(Arc::new))
-                .map_err(RbPolarsErr::from)?;
-            BatchedReader::MMap(reader)
-        };
+        let reader = reader
+            .batched(overwrite_dtype.map(Arc::new))
+            .map_err(RbPolarsErr::from)?;
 
         Ok(RbBatchedCsv {
-            reader: RefCell::new(reader),
+            reader: RefCell::new(Mutex::new(reader)),
         })
     }
 
     pub fn next_batches(&self, n: usize) -> RbResult<Option<RArray>> {
-        let batches = match &mut *self.reader.borrow_mut() {
-            BatchedReader::MMap(reader) => reader.next_batches(n),
-            BatchedReader::Read(reader) => reader.next_batches(n),
-        }
-        .map_err(RbPolarsErr::from)?;
+        let reader = &self.reader;
+        let batches = reader
+            .borrow()
+            .lock()
+            .map_err(|e| RbPolarsErr::other(e.to_string()))?
+            .next_batches(n)
+            .map_err(RbPolarsErr::from)?;
 
         Ok(batches.map(|batches| RArray::from_iter(batches.into_iter().map(RbDataFrame::from))))
     }

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(),