polars-df 0.1.0 → 0.1.2

data/ext/polars/src/dataframe.rs

@@ -1,5 +1,6 @@
 use magnus::{r_hash::ForEach, Error, RArray, RHash, RString, Value};
 use polars::io::mmap::ReaderBytes;
+use polars::io::RowCount;
 use polars::prelude::*;
 use std::cell::RefCell;
 use std::fs::File;
@@ -7,8 +8,9 @@ use std::io::{BufReader, BufWriter, Cursor};
 use std::ops::Deref;
 use std::path::PathBuf;
 
-use crate::conversion::parse_parquet_compression;
+use crate::conversion::*;
 use crate::file::{get_file_like, get_mmap_bytes_reader};
+use crate::series::to_rbseries_collection;
 use crate::{series, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
 
 #[magnus::wrap(class = "Polars::RbDataFrame")]
@@ -38,10 +40,98 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
-    pub fn read_csv(rb_f: Value, has_header: bool) -> RbResult<Self> {
+    pub fn estimated_size(&self) -> usize {
+        self.df.borrow().estimated_size()
+    }
+
+    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()?;
+        // 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 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()?;
+        // end arguments
+
+        let null_values = null_values.map(|w| w.0);
+        let comment_char = comment_char.map(|s| s.as_bytes()[0]);
+        let eol_char = eol_char.as_bytes()[0];
+
+        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+
+        let quote_char = if let Some(s) = quote_char {
+            if s.is_empty() {
+                None
+            } else {
+                Some(s.as_bytes()[0])
+            }
+        } else {
+            None
+        };
+
+        let overwrite_dtype = overwrite_dtype.map(|overwrite_dtype| {
+            let fields = overwrite_dtype.iter().map(|(name, dtype)| {
+                let dtype = dtype.0.clone();
+                Field::new(name, dtype)
+            });
+            Schema::from(fields)
+        });
+
+        let overwrite_dtype_slice = overwrite_dtype_slice.map(|overwrite_dtype| {
+            overwrite_dtype
+                .iter()
+                .map(|dt| dt.0.clone())
+                .collect::<Vec<_>>()
+        });
+
         let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
         let df = CsvReader::new(mmap_bytes_r)
+            .infer_schema(infer_schema_length)
             .has_header(has_header)
+            .with_n_rows(n_rows)
+            .with_delimiter(sep.as_bytes()[0])
+            .with_skip_rows(skip_rows)
+            .with_ignore_parser_errors(ignore_errors)
+            .with_projection(projection)
+            .with_rechunk(rechunk)
+            .with_chunk_size(chunk_size)
+            .with_encoding(encoding.0)
+            .with_columns(columns)
+            .with_n_threads(n_threads)
+            .with_path(path)
+            .with_dtypes(overwrite_dtype.as_ref())
+            .with_dtypes_slice(overwrite_dtype_slice.as_deref())
+            .low_memory(low_memory)
+            .with_comment_char(comment_char)
+            .with_null_values(null_values)
+            .with_parse_dates(parse_dates)
+            .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)
             .finish()
             .map_err(RbPolarsErr::from)?;
         Ok(df.into())
@@ -56,6 +146,27 @@ impl RbDataFrame {
             .map(|v| v.into())
     }
 
+    pub fn read_ipc(
+        rb_f: Value,
+        columns: Option<Vec<String>>,
+        projection: Option<Vec<usize>>,
+        n_rows: Option<usize>,
+        row_count: Option<(String, IdxSize)>,
+        memory_map: bool,
+    ) -> RbResult<Self> {
+        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+        let df = IpcReader::new(mmap_bytes_r)
+            .with_projection(projection)
+            .with_columns(columns)
+            .with_n_rows(n_rows)
+            .with_row_count(row_count)
+            .memory_mapped(memory_map)
+            .finish()
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
     pub fn read_json(rb_f: Value) -> RbResult<Self> {
         // memmap the file first
         let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
@@ -180,6 +291,28 @@ impl RbDataFrame {
         Ok(())
     }
 
+    pub fn write_ipc(
+        &self,
+        rb_f: Value,
+        compression: Wrap<Option<IpcCompression>>,
+    ) -> RbResult<()> {
+        if let Ok(s) = rb_f.try_convert::<String>() {
+            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)?;
+
+            IpcWriter::new(&mut buf)
+                .with_compression(compression.0)
+                .finish(&mut self.df.borrow_mut())
+                .map_err(RbPolarsErr::from)?;
+        }
+        Ok(())
+    }
+
     pub fn write_parquet(
         &self,
         rb_f: Value,
@@ -213,6 +346,11 @@ impl RbDataFrame {
         format!("{}", self.df.borrow())
     }
 
+    pub fn get_columns(&self) -> Vec<RbSeries> {
+        let cols = self.df.borrow().get_columns().clone();
+        to_rbseries_collection(cols)
+    }
+
     pub fn columns(&self) -> Vec<String> {
         self.df
             .borrow()
@@ -222,14 +360,27 @@ impl RbDataFrame {
             .collect()
     }
 
-    pub fn dtypes(&self) -> Vec<String> {
+    pub fn set_column_names(&self, names: Vec<String>) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .set_column_names(&names)
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn dtypes(&self) -> Vec<Value> {
         self.df
             .borrow()
             .iter()
-            .map(|s| s.dtype().to_string())
+            .map(|s| Wrap(s.dtype().clone()).into())
             .collect()
     }
 
+    pub fn n_chunks(&self) -> RbResult<usize> {
+        let n = self.df.borrow().n_chunks().map_err(RbPolarsErr::from)?;
+        Ok(n)
+    }
+
     pub fn shape(&self) -> (usize, usize) {
         self.df.borrow().shape()
     }
@@ -258,6 +409,28 @@ impl RbDataFrame {
             .map_err(RbPolarsErr::from)
     }
 
+    pub fn select(&self, selection: Vec<String>) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .select(selection)
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
+    pub fn take(&self, indices: Vec<IdxSize>) -> RbResult<Self> {
+        let indices = IdxCa::from_vec("", indices);
+        let df = self.df.borrow().take(&indices).map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
+    pub fn take_with_series(&self, indices: &RbSeries) -> RbResult<Self> {
+        let binding = indices.series.borrow();
+        let idx = binding.idx().map_err(RbPolarsErr::from)?;
+        let df = self.df.borrow().take(idx).map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
     pub fn sort(&self, by_column: String, reverse: bool, nulls_last: bool) -> RbResult<Self> {
         let df = self
             .df
@@ -273,6 +446,38 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
+    pub fn replace(&self, column: String, new_col: &RbSeries) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .replace(&column, new_col.series.borrow().clone())
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn replace_at_idx(&self, index: usize, new_col: &RbSeries) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .replace_at_idx(index, new_col.series.borrow().clone())
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn insert_at_idx(&self, index: usize, new_col: &RbSeries) -> RbResult<()> {
+        self.df
+            .borrow_mut()
+            .insert_at_idx(index, new_col.series.borrow().clone())
+            .map_err(RbPolarsErr::from)?;
+        Ok(())
+    }
+
+    pub fn slice(&self, offset: usize, length: Option<usize>) -> Self {
+        let df = self.df.borrow().slice(
+            offset as i64,
+            length.unwrap_or_else(|| self.df.borrow().height()),
+        );
+        df.into()
+    }
+
     pub fn head(&self, length: Option<usize>) -> Self {
         self.df.borrow().head(length).into()
     }
@@ -281,6 +486,20 @@ impl RbDataFrame {
         self.df.borrow().tail(length).into()
     }
 
+    pub fn is_unique(&self) -> RbResult<RbSeries> {
+        let mask = self.df.borrow().is_unique().map_err(RbPolarsErr::from)?;
+        Ok(mask.into_series().into())
+    }
+
+    pub fn is_duplicated(&self) -> RbResult<RbSeries> {
+        let mask = self
+            .df
+            .borrow()
+            .is_duplicated()
+            .map_err(RbPolarsErr::from)?;
+        Ok(mask.into_series().into())
+    }
+
     pub fn frame_equal(&self, other: &RbDataFrame, null_equal: bool) -> bool {
         if null_equal {
             self.df.borrow().frame_equal_missing(&other.df.borrow())
@@ -289,16 +508,202 @@ impl RbDataFrame {
         }
     }
 
+    pub fn with_row_count(&self, name: String, offset: Option<IdxSize>) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .with_row_count(&name, offset)
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn clone(&self) -> Self {
+        RbDataFrame::new(self.df.borrow().clone())
+    }
+
+    pub fn melt(
+        &self,
+        id_vars: Vec<String>,
+        value_vars: Vec<String>,
+        value_name: Option<String>,
+        variable_name: Option<String>,
+    ) -> RbResult<Self> {
+        let args = MeltArgs {
+            id_vars,
+            value_vars,
+            value_name,
+            variable_name,
+        };
+
+        let df = self.df.borrow().melt2(args).map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
+    pub fn partition_by(&self, groups: Vec<String>, stable: bool) -> RbResult<Vec<Self>> {
+        let out = if stable {
+            self.df.borrow().partition_by_stable(groups)
+        } else {
+            self.df.borrow().partition_by(groups)
+        }
+        .map_err(RbPolarsErr::from)?;
+        Ok(out.into_iter().map(RbDataFrame::new).collect())
+    }
+
+    pub fn shift(&self, periods: i64) -> Self {
+        self.df.borrow().shift(periods).into()
+    }
+
+    pub fn unique(
+        &self,
+        maintain_order: bool,
+        subset: Option<Vec<String>>,
+        keep: Wrap<UniqueKeepStrategy>,
+    ) -> RbResult<Self> {
+        let subset = subset.as_ref().map(|v| v.as_ref());
+        let df = match maintain_order {
+            true => self.df.borrow().unique_stable(subset, keep.0),
+            false => self.df.borrow().unique(subset, keep.0),
+        }
+        .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
     pub fn lazy(&self) -> RbLazyFrame {
         self.df.borrow().clone().lazy().into()
     }
 
+    pub fn max(&self) -> Self {
+        self.df.borrow().max().into()
+    }
+
+    pub fn min(&self) -> Self {
+        self.df.borrow().min().into()
+    }
+
+    pub fn sum(&self) -> Self {
+        self.df.borrow().sum().into()
+    }
+
     pub fn mean(&self) -> Self {
         self.df.borrow().mean().into()
     }
 
+    pub fn std(&self, ddof: u8) -> Self {
+        self.df.borrow().std(ddof).into()
+    }
+
+    pub fn var(&self, ddof: u8) -> Self {
+        self.df.borrow().var(ddof).into()
+    }
+
+    pub fn median(&self) -> Self {
+        self.df.borrow().median().into()
+    }
+
+    pub fn hmean(&self, null_strategy: Wrap<NullStrategy>) -> RbResult<Option<RbSeries>> {
+        let s = self
+            .df
+            .borrow()
+            .hmean(null_strategy.0)
+            .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)?;
+        Ok(s.map(|s| s.into()))
+    }
+
+    pub fn hmin(&self) -> RbResult<Option<RbSeries>> {
+        let s = self.df.borrow().hmin().map_err(RbPolarsErr::from)?;
+        Ok(s.map(|s| s.into()))
+    }
+
+    pub fn hsum(&self, null_strategy: Wrap<NullStrategy>) -> RbResult<Option<RbSeries>> {
+        let s = self
+            .df
+            .borrow()
+            .hsum(null_strategy.0)
+            .map_err(RbPolarsErr::from)?;
+        Ok(s.map(|s| s.into()))
+    }
+
+    pub fn quantile(
+        &self,
+        quantile: f64,
+        interpolation: Wrap<QuantileInterpolOptions>,
+    ) -> RbResult<Self> {
+        let df = self
+            .df
+            .borrow()
+            .quantile(quantile, interpolation.0)
+            .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
+    pub fn to_dummies(&self, columns: Option<Vec<String>>) -> RbResult<Self> {
+        let df = match columns {
+            Some(cols) => self
+                .df
+                .borrow()
+                .columns_to_dummies(cols.iter().map(|x| x as &str).collect()),
+            None => self.df.borrow().to_dummies(),
+        }
+        .map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
+
     pub fn null_count(&self) -> Self {
         let df = self.df.borrow().null_count();
         df.into()
     }
+
+    pub fn shrink_to_fit(&self) {
+        self.df.borrow_mut().shrink_to_fit();
+    }
+
+    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 upsample(
+        &self,
+        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),
+            )
+        } else {
+            self.df.borrow().upsample(
+                by,
+                &index_column,
+                Duration::parse(&every),
+                Duration::parse(&offset),
+            )
+        };
+        let out = out.map_err(RbPolarsErr::from)?;
+        Ok(out.into())
+    }
+
+    pub fn unnest(&self, names: Vec<String>) -> RbResult<Self> {
+        let df = self.df.borrow().unnest(names).map_err(RbPolarsErr::from)?;
+        Ok(df.into())
+    }
 }

data/ext/polars/src/error.rs

@@ -1,5 +1,6 @@
 use magnus::exception::arg_error;
 use magnus::Error;
+use polars::error::ArrowError;
 use polars::prelude::PolarsError;
 
 pub struct RbPolarsErr {}
@@ -10,6 +11,14 @@ impl RbPolarsErr {
         Error::runtime_error(e.to_string())
     }
 
+    pub fn arrow(e: ArrowError) -> Error {
+        Error::runtime_error(e.to_string())
+    }
+
+    pub fn io(e: std::io::Error) -> Error {
+        Error::runtime_error(e.to_string())
+    }
+
     pub fn other(message: String) -> Error {
         Error::runtime_error(message)
     }

data/ext/polars/src/file.rs

@@ -1,18 +1,19 @@
 use magnus::{Error, RString, Value};
 use polars::io::mmap::MmapBytesReader;
-use std::fs::{File, OpenOptions};
+use std::fs::File;
 use std::io::Cursor;
 use std::path::PathBuf;
 
 use crate::RbResult;
 
 pub fn get_file_like(f: Value, truncate: bool) -> RbResult<File> {
-    OpenOptions::new()
-        .write(true)
-        .create(true)
-        .truncate(truncate)
-        .open(f.try_convert::<PathBuf>()?)
-        .map_err(|e| Error::runtime_error(e.to_string()))
+    let str_slice = f.try_convert::<PathBuf>()?;
+    let f = if truncate {
+        File::create(str_slice).map_err(|e| Error::runtime_error(e.to_string()))?
+    } else {
+        File::open(str_slice).map_err(|e| Error::runtime_error(e.to_string()))?
+    };
+    Ok(f)
 }
 
 pub fn get_mmap_bytes_reader(rb_f: Value) -> RbResult<Box<dyn MmapBytesReader>> {

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

@@ -0,0 +1,7 @@
+use magnus::Value;
+use polars::error::PolarsResult;
+use polars::series::Series;
+
+pub fn binary_lambda(_lambda: Value, _a: Series, _b: Series) -> PolarsResult<Series> {
+    todo!();
+}