polars-df 0.5.0 → 0.7.0

data/ext/polars/src/dataframe.rs

@@ -1,4 +1,7 @@
-use magnus::{r_hash::ForEach, IntoValue, RArray, RHash, RString, Value};
+use either::Either;
+use magnus::{
+    prelude::*, r_hash::ForEach, typed_data::Obj, IntoValue, RArray, RHash, RString, Value,
+};
 use polars::frame::row::{rows_to_schema_supertypes, Row};
 use polars::frame::NullStrategy;
 use polars::io::avro::AvroCompression;
@@ -11,12 +14,12 @@ use std::cell::RefCell;
 use std::io::{BufWriter, Cursor};
 use std::ops::Deref;
 
-use crate::apply::dataframe::{
+use crate::conversion::*;
+use crate::file::{get_file_like, get_mmap_bytes_reader};
+use crate::map::dataframe::{
     apply_lambda_unknown, apply_lambda_with_bool_out_type, apply_lambda_with_primitive_out_type,
     apply_lambda_with_utf8_out_type,
 };
-use crate::conversion::*;
-use crate::file::{get_file_like, get_mmap_bytes_reader};
 use crate::rb_modules;
 use crate::series::{to_rbseries_collection, to_series_collection};
 use crate::{RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
@@ -53,7 +56,7 @@ impl RbDataFrame {
         // replace inferred nulls with boolean
         let fields = schema.iter_fields().map(|mut fld| match fld.data_type() {
             DataType::Null => {
-                fld.coerce(DataType::Boolean);
+                // fld.coerce(DataType::Boolean);
                 fld
             }
             DataType::Decimal(_, _) => {
@@ -86,7 +89,7 @@ impl RbDataFrame {
     pub fn init(columns: RArray) -> RbResult<Self> {
         let mut cols = Vec::new();
         for i in columns.each() {
-            cols.push(i?.try_convert::<&RbSeries>()?.series.borrow().clone());
+            cols.push(<&RbSeries>::try_convert(i?)?.series.borrow().clone());
         }
         let df = DataFrame::new(cols).map_err(RbPolarsErr::from)?;
         Ok(RbDataFrame::new(df))
@@ -99,32 +102,32 @@ impl RbDataFrame {
     pub fn read_csv(arguments: &[Value]) -> RbResult<Self> {
         // start arguments
         // this pattern is needed for more than 16
-        let rb_f: Value = arguments[0].try_convert()?;
-        let infer_schema_length: Option<usize> = arguments[1].try_convert()?;
-        let chunk_size: usize = arguments[2].try_convert()?;
-        let has_header: bool = arguments[3].try_convert()?;
-        let ignore_errors: bool = arguments[4].try_convert()?;
-        let n_rows: Option<usize> = arguments[5].try_convert()?;
-        let skip_rows: usize = arguments[6].try_convert()?;
-        let projection: Option<Vec<usize>> = arguments[7].try_convert()?;
-        let sep: String = arguments[8].try_convert()?;
-        let rechunk: bool = arguments[9].try_convert()?;
-        let columns: Option<Vec<String>> = arguments[10].try_convert()?;
-        let encoding: Wrap<CsvEncoding> = arguments[11].try_convert()?;
-        let n_threads: Option<usize> = arguments[12].try_convert()?;
-        let path: Option<String> = arguments[13].try_convert()?;
-        let overwrite_dtype: Option<Vec<(String, Wrap<DataType>)>> = arguments[14].try_convert()?;
+        let rb_f = arguments[0];
+        let infer_schema_length = Option::<usize>::try_convert(arguments[1])?;
+        let chunk_size = usize::try_convert(arguments[2])?;
+        let has_header = bool::try_convert(arguments[3])?;
+        let ignore_errors = bool::try_convert(arguments[4])?;
+        let n_rows = Option::<usize>::try_convert(arguments[5])?;
+        let skip_rows = usize::try_convert(arguments[6])?;
+        let projection = Option::<Vec<usize>>::try_convert(arguments[7])?;
+        let separator = String::try_convert(arguments[8])?;
+        let rechunk = bool::try_convert(arguments[9])?;
+        let columns = Option::<Vec<String>>::try_convert(arguments[10])?;
+        let encoding = Wrap::<CsvEncoding>::try_convert(arguments[11])?;
+        let n_threads = Option::<usize>::try_convert(arguments[12])?;
+        let path = Option::<String>::try_convert(arguments[13])?;
+        let overwrite_dtype = Option::<Vec<(String, Wrap<DataType>)>>::try_convert(arguments[14])?;
         // TODO fix
-        let overwrite_dtype_slice: Option<Vec<Wrap<DataType>>> = None; // arguments[15].try_convert()?;
-        let low_memory: bool = arguments[16].try_convert()?;
-        let comment_char: Option<String> = arguments[17].try_convert()?;
-        let quote_char: Option<String> = arguments[18].try_convert()?;
-        let null_values: Option<Wrap<NullValues>> = arguments[19].try_convert()?;
-        let try_parse_dates: bool = arguments[20].try_convert()?;
-        let skip_rows_after_header: usize = arguments[21].try_convert()?;
-        let row_count: Option<(String, IdxSize)> = arguments[22].try_convert()?;
-        let sample_size: usize = arguments[23].try_convert()?;
-        let eol_char: String = arguments[24].try_convert()?;
+        let overwrite_dtype_slice = Option::<Vec<Wrap<DataType>>>::None; // Option::<Vec<Wrap<DataType>>>::try_convert(arguments[15])?;
+        let low_memory = bool::try_convert(arguments[16])?;
+        let comment_char = Option::<String>::try_convert(arguments[17])?;
+        let quote_char = Option::<String>::try_convert(arguments[18])?;
+        let null_values = Option::<Wrap<NullValues>>::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_count = Option::<(String, IdxSize)>::try_convert(arguments[22])?;
+        let sample_size = usize::try_convert(arguments[23])?;
+        let eol_char = String::try_convert(arguments[24])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
@@ -165,7 +168,7 @@ impl RbDataFrame {
             .infer_schema(infer_schema_length)
             .has_header(has_header)
             .with_n_rows(n_rows)
-            .with_delimiter(sep.as_bytes()[0])
+            .with_separator(separator.as_bytes()[0])
             .with_skip_rows(skip_rows)
             .with_ignore_errors(ignore_errors)
             .with_projection(projection)
@@ -265,7 +268,7 @@ impl RbDataFrame {
     ) -> RbResult<()> {
         use polars::io::avro::AvroWriter;
 
-        if let Ok(s) = rb_f.try_convert::<String>() {
+        if let Ok(s) = String::try_convert(rb_f) {
             let f = std::fs::File::create(s).unwrap();
             AvroWriter::new(f)
                 .with_compression(compression.0)
@@ -341,6 +344,27 @@ impl RbDataFrame {
         Ok(())
     }
 
+    pub fn read_rows(
+        rb_rows: RArray,
+        infer_schema_length: Option<usize>,
+        schema_overwrite: Option<Wrap<Schema>>,
+    ) -> RbResult<Self> {
+        let mut rows = Vec::with_capacity(rb_rows.len());
+        for v in rb_rows.each() {
+            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);
+            }
+            rows.push(Row(row));
+        }
+        Self::finish_from_rows(
+            rows,
+            infer_schema_length,
+            schema_overwrite.map(|wrap| wrap.0),
+        )
+    }
+
     pub fn read_hashes(
         dicts: Value,
         infer_schema_length: Option<usize>,
@@ -395,9 +419,9 @@ impl RbDataFrame {
     pub fn write_csv(
         &self,
         rb_f: Value,
-        has_header: bool,
-        sep: u8,
-        quote: u8,
+        include_header: bool,
+        separator: u8,
+        quote_char: u8,
         batch_size: usize,
         datetime_format: Option<String>,
         date_format: Option<String>,
@@ -407,13 +431,13 @@ impl RbDataFrame {
     ) -> RbResult<()> {
         let null = null_value.unwrap_or_default();
 
-        if let Ok(s) = rb_f.try_convert::<String>() {
+        if let Ok(s) = String::try_convert(rb_f) {
             let f = std::fs::File::create(s).unwrap();
             // no need for a buffered writer, because the csv writer does internal buffering
             CsvWriter::new(f)
-                .has_header(has_header)
-                .with_delimiter(sep)
-                .with_quoting_char(quote)
+                .include_header(include_header)
+                .with_separator(separator)
+                .with_quote_char(quote_char)
                 .with_batch_size(batch_size)
                 .with_datetime_format(datetime_format)
                 .with_date_format(date_format)
@@ -425,9 +449,9 @@ impl RbDataFrame {
         } else {
             let mut buf = Cursor::new(Vec::new());
             CsvWriter::new(&mut buf)
-                .has_header(has_header)
-                .with_delimiter(sep)
-                .with_quoting_char(quote)
+                .include_header(include_header)
+                .with_separator(separator)
+                .with_quote_char(quote_char)
                 .with_batch_size(batch_size)
                 .with_datetime_format(datetime_format)
                 .with_date_format(date_format)
@@ -449,19 +473,21 @@ impl RbDataFrame {
         rb_f: Value,
         compression: Wrap<Option<IpcCompression>>,
     ) -> RbResult<()> {
-        if let Ok(s) = rb_f.try_convert::<String>() {
+        if let Ok(s) = String::try_convert(rb_f) {
             let f = std::fs::File::create(s).unwrap();
             IpcWriter::new(f)
                 .with_compression(compression.0)
                 .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(())
     }
@@ -485,7 +511,7 @@ impl RbDataFrame {
                     _ => Wrap(s.get(idx).unwrap()).into_value(),
                 }),
         )
-        .into()
+        .as_value()
     }
 
     pub fn row_tuples(&self) -> Value {
@@ -505,7 +531,7 @@ impl RbDataFrame {
                     }),
             )
         }))
-        .into()
+        .as_value()
     }
 
     pub fn to_numo(&self) -> Option<Value> {
@@ -535,7 +561,7 @@ impl RbDataFrame {
     ) -> RbResult<()> {
         let compression = parse_parquet_compression(&compression, compression_level)?;
 
-        if let Ok(s) = rb_f.try_convert::<String>() {
+        if let Ok(s) = String::try_convert(rb_f) {
             let f = std::fs::File::create(s).unwrap();
             ParquetWriter::new(f)
                 .with_compression(compression)
@@ -602,7 +628,7 @@ impl RbDataFrame {
 
     pub fn sample_n(
         &self,
-        n: usize,
+        n: &RbSeries,
         with_replacement: bool,
         shuffle: bool,
         seed: Option<u64>,
@@ -610,14 +636,14 @@ impl RbDataFrame {
         let df = self
             .df
             .borrow()
-            .sample_n(n, with_replacement, shuffle, seed)
+            .sample_n(&n.series.borrow(), with_replacement, shuffle, seed)
             .map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }
 
     pub fn sample_frac(
         &self,
-        frac: f64,
+        frac: &RbSeries,
         with_replacement: bool,
         shuffle: bool,
         seed: Option<u64>,
@@ -625,7 +651,7 @@ impl RbDataFrame {
         let df = self
             .df
             .borrow()
-            .sample_frac(frac, with_replacement, shuffle, seed)
+            .sample_frac(&frac.series.borrow(), with_replacement, shuffle, seed)
             .map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }
@@ -794,22 +820,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 +943,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)))
@@ -979,30 +994,48 @@ impl RbDataFrame {
         self.df.borrow().median().into()
     }
 
-    pub fn hmean(&self, null_strategy: Wrap<NullStrategy>) -> RbResult<Option<RbSeries>> {
+    pub fn max_horizontal(&self) -> RbResult<Option<RbSeries>> {
         let s = self
             .df
             .borrow()
-            .hmean(null_strategy.0)
+            .max_horizontal()
             .map_err(RbPolarsErr::from)?;
         Ok(s.map(|s| s.into()))
     }
 
-    pub fn hmax(&self) -> RbResult<Option<RbSeries>> {
-        let s = self.df.borrow().hmax().map_err(RbPolarsErr::from)?;
+    pub fn min_horizontal(&self) -> RbResult<Option<RbSeries>> {
+        let s = self
+            .df
+            .borrow()
+            .min_horizontal()
+            .map_err(RbPolarsErr::from)?;
         Ok(s.map(|s| s.into()))
     }
 
-    pub fn hmin(&self) -> RbResult<Option<RbSeries>> {
-        let s = self.df.borrow().hmin().map_err(RbPolarsErr::from)?;
+    pub fn sum_horizontal(&self, ignore_nulls: bool) -> RbResult<Option<RbSeries>> {
+        let null_strategy = if ignore_nulls {
+            NullStrategy::Ignore
+        } else {
+            NullStrategy::Propagate
+        };
+        let s = self
+            .df
+            .borrow()
+            .sum_horizontal(null_strategy)
+            .map_err(RbPolarsErr::from)?;
         Ok(s.map(|s| s.into()))
     }
 
-    pub fn hsum(&self, null_strategy: Wrap<NullStrategy>) -> RbResult<Option<RbSeries>> {
+    pub fn mean_horizontal(&self, ignore_nulls: bool) -> RbResult<Option<RbSeries>> {
+        let null_strategy = if ignore_nulls {
+            NullStrategy::Ignore
+        } else {
+            NullStrategy::Propagate
+        };
         let s = self
             .df
             .borrow()
-            .hsum(null_strategy.0)
+            .mean_horizontal(null_strategy)
             .map_err(RbPolarsErr::from)?;
         Ok(s.map(|s| s.into()))
     }
@@ -1024,13 +1057,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())
@@ -1092,7 +1130,7 @@ impl RbDataFrame {
             _ => return apply_lambda_unknown(df, lambda, inference_size),
         };
 
-        Ok((RbSeries::from(out).into(), false))
+        Ok((Obj::wrap(RbSeries::from(out)).as_value(), false))
     }
 
     pub fn shrink_to_fit(&self) {
@@ -1109,17 +1147,20 @@ impl RbDataFrame {
         Ok(hash.into_series().into())
     }
 
-    pub fn transpose(&self, include_header: bool, names: String) -> RbResult<Self> {
-        let mut df = self.df.borrow().transpose().map_err(RbPolarsErr::from)?;
-        if include_header {
-            let s = Utf8Chunked::from_iter_values(
-                &names,
-                self.df.borrow().get_columns().iter().map(|s| s.name()),
-            )
-            .into_series();
-            df.insert_at_idx(0, s).unwrap();
-        }
-        Ok(df.into())
+    pub fn transpose(&self, keep_names_as: Option<String>, column_names: Value) -> RbResult<Self> {
+        let new_col_names = if let Ok(name) = <Vec<String>>::try_convert(column_names) {
+            Some(Either::Right(name))
+        } else if let Ok(name) = String::try_convert(column_names) {
+            Some(Either::Left(name))
+        } else {
+            None
+        };
+        Ok(self
+            .df
+            .borrow()
+            .transpose(keep_names_as.as_deref(), new_col_names)
+            .map_err(RbPolarsErr::from)?
+            .into())
     }
 
     pub fn upsample(

data/ext/polars/src/error.rs

@@ -1,6 +1,5 @@
 use magnus::exception;
 use magnus::Error;
-use polars::error::ArrowError;
 use polars::prelude::PolarsError;
 
 pub struct RbPolarsErr {}
@@ -11,10 +10,6 @@ impl RbPolarsErr {
         Error::new(exception::runtime_error(), e.to_string())
     }
 
-    pub fn arrow(e: ArrowError) -> Error {
-        Error::new(exception::runtime_error(), e.to_string())
-    }
-
     pub fn io(e: std::io::Error) -> Error {
         Error::new(exception::runtime_error(), e.to_string())
     }
@@ -24,6 +19,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()
+    }
+}

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

@@ -3,16 +3,28 @@ use polars::prelude::*;
 use crate::RbExpr;
 
 impl RbExpr {
-    pub fn bin_contains(&self, lit: Vec<u8>) -> Self {
-        self.inner.clone().binary().contains_literal(lit).into()
+    pub fn bin_contains(&self, lit: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .binary()
+            .contains_literal(lit.inner.clone())
+            .into()
     }
 
-    pub fn bin_ends_with(&self, sub: Vec<u8>) -> Self {
-        self.inner.clone().binary().ends_with(sub).into()
+    pub fn bin_ends_with(&self, sub: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .binary()
+            .ends_with(sub.inner.clone())
+            .into()
     }
 
-    pub fn bin_starts_with(&self, sub: Vec<u8>) -> Self {
-        self.inner.clone().binary().starts_with(sub).into()
+    pub fn bin_starts_with(&self, sub: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .binary()
+            .starts_with(sub.inner.clone())
+            .into()
     }
 
     pub fn bin_hex_decode(&self, strict: bool) -> Self {

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

@@ -8,9 +8,8 @@ impl RbExpr {
         self.inner.clone().dt().to_string(&format).into()
     }
 
-    pub fn dt_offset_by(&self, by: String) -> Self {
-        let by = Duration::parse(&by);
-        self.inner.clone().dt().offset_by(by).into()
+    pub fn dt_offset_by(&self, by: &RbExpr) -> Self {
+        self.inner.clone().dt().offset_by(by.inner.clone()).into()
     }
 
     pub fn dt_epoch_seconds(&self) -> Self {
@@ -38,21 +37,20 @@ impl RbExpr {
         self.inner.clone().dt().cast_time_unit(tu.0).into()
     }
 
-    pub fn dt_replace_time_zone(&self, tz: Option<String>, use_earliest: Option<bool>) -> Self {
+    pub fn dt_replace_time_zone(&self, time_zone: Option<String>, ambiguous: &Self) -> Self {
         self.inner
             .clone()
             .dt()
-            .replace_time_zone(tz, use_earliest)
+            .replace_time_zone(time_zone, ambiguous.inner.clone())
             .into()
     }
 
-    #[allow(deprecated)]
-    pub fn dt_tz_localize(&self, tz: String) -> Self {
-        self.inner.clone().dt().tz_localize(tz).into()
-    }
-
-    pub fn dt_truncate(&self, every: String, offset: String) -> Self {
-        self.inner.clone().dt().truncate(&every, &offset).into()
+    pub fn dt_truncate(&self, every: &Self, offset: String) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .truncate(every.inner.clone(), offset)
+            .into()
     }
 
     pub fn dt_month_start(&self) -> Self {