polars-df 0.1.1 → 0.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2a39332a375d211cf0e04ea95cdb4fe315c843b3df474ee3fef135ec63eb0516
-  data.tar.gz: 41881b2aa8da2bb2850fbcaba11ec1409204ec99b7d36d1470a7536684b02509
+  metadata.gz: 510c7a761553fb3a49919add842d520da1518c5df6dd37f93af338e928c7a207
+  data.tar.gz: 3e724d3c2553bb6b4a587f056a04977990b85103ea9ae231166746407e1c0b1e
 SHA512:
-  metadata.gz: 17ba7456808346b025d52d14d5fa6d91dc0d18385d0ee796c1b821eb6272875e15927e34da8f41c1872e4b158bfd32165ce1ee4adc5cebf4474fc7e86ad98e31
-  data.tar.gz: 4c3ef1aa3ed5c6e0620d555bab73d20e8b7cd495e26c97f4ddd41f14ed8d3883331cff36158fbe4e1dbc3b86825fa83dfa4b832f901db35a73a9c021b8dd955c
+  metadata.gz: 59391ecd7a3d14372ad24693240ebeae4d83c59d55b37876fc6f216dd576897113607f7b5f6f6a3d7a5d75f96def72b1f2659f3de71e6a975a2753719879fa33
+  data.tar.gz: 5a2e8cdf3aed01f16823be62067b0ec01e4e0dc3bfa4dc27929dfbe42d7c92e62f6e7a09bc78533d220a3bfc9890f9f0fb069482a338fbd0b55603aa42fddc57

data/CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.1.2 (2022-11-25)
+
+- Added more methods
+
 ## 0.1.1 (2022-11-23)
 
 - Added more methods

data/Cargo.lock

@@ -1160,7 +1160,7 @@ dependencies = [
 
 [[package]]
 name = "polars"
-version = "0.1.0"
+version = "0.1.2"
 dependencies = [
  "magnus",
  "polars 0.25.1",

data/ext/polars/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "polars"
-version = "0.1.1"
+version = "0.1.2"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
 publish = false

data/ext/polars/src/batched_csv.rs

@@ -0,0 +1,120 @@
+use magnus::Value;
+use polars::io::mmap::MmapBytesReader;
+use polars::io::RowCount;
+use polars::prelude::read_impl::OwnedBatchedCsvReader;
+use polars::prelude::*;
+use std::cell::RefCell;
+use std::path::PathBuf;
+
+use crate::conversion::*;
+use crate::{RbDataFrame, RbPolarsErr, RbResult};
+
+#[magnus::wrap(class = "Polars::RbBatchedCsv")]
+pub struct RbBatchedCsv {
+    pub reader: RefCell<OwnedBatchedCsvReader>,
+}
+
+impl RbBatchedCsv {
+    pub fn new(arguments: &[Value]) -> RbResult<Self> {
+        // start arguments
+        // this pattern is needed for more than 16
+        let infer_schema_length: Option<usize> = arguments[0].try_convert()?;
+        let chunk_size: usize = arguments[1].try_convert()?;
+        let has_header: bool = arguments[2].try_convert()?;
+        let ignore_errors: bool = arguments[3].try_convert()?;
+        let n_rows: Option<usize> = arguments[4].try_convert()?;
+        let skip_rows: usize = arguments[5].try_convert()?;
+        let projection: Option<Vec<usize>> = arguments[6].try_convert()?;
+        let sep: String = arguments[7].try_convert()?;
+        let rechunk: bool = arguments[8].try_convert()?;
+        let columns: Option<Vec<String>> = arguments[9].try_convert()?;
+        let encoding: Wrap<CsvEncoding> = arguments[10].try_convert()?;
+        let n_threads: Option<usize> = arguments[11].try_convert()?;
+        let path: PathBuf = arguments[12].try_convert()?;
+        let overwrite_dtype: Option<Vec<(String, Wrap<DataType>)>> = arguments[13].try_convert()?;
+        // TODO fix
+        let overwrite_dtype_slice: Option<Vec<Wrap<DataType>>> = None; // arguments[14].try_convert()?;
+        let low_memory: bool = arguments[15].try_convert()?;
+        let comment_char: Option<String> = arguments[16].try_convert()?;
+        let quote_char: Option<String> = arguments[17].try_convert()?;
+        let null_values: Option<Wrap<NullValues>> = arguments[18].try_convert()?;
+        let parse_dates: bool = arguments[19].try_convert()?;
+        let skip_rows_after_header: usize = arguments[20].try_convert()?;
+        let row_count: Option<(String, IdxSize)> = arguments[21].try_convert()?;
+        let sample_size: usize = arguments[22].try_convert()?;
+        let eol_char: String = arguments[23].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 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)
+            .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_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)
+            .batched(overwrite_dtype.map(Arc::new))
+            .map_err(RbPolarsErr::from)?;
+
+        Ok(RbBatchedCsv {
+            reader: RefCell::new(reader),
+        })
+    }
+
+    pub fn next_batches(&self, n: usize) -> RbResult<Option<Vec<RbDataFrame>>> {
+        let batches = self
+            .reader
+            .borrow_mut()
+            .next_batches(n)
+            .map_err(RbPolarsErr::from)?;
+        Ok(batches.map(|batches| batches.into_iter().map(|out| out.1.into()).collect()))
+    }
+}

data/ext/polars/src/conversion.rs

@@ -1,11 +1,11 @@
-use magnus::{TryConvert, Value, QNIL};
+use magnus::{RArray, Symbol, TryConvert, Value, QNIL};
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
 use polars::datatypes::AnyValue;
 use polars::frame::DataFrame;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 
-use crate::{RbDataFrame, RbPolarsErr, RbResult, RbValueError};
+use crate::{RbDataFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
 pub struct Wrap<T>(pub T);
 
@@ -15,14 +15,57 @@ impl<T> From<T> for Wrap<T> {
     }
 }
 
+pub fn get_rbseq(obj: Value) -> RbResult<(RArray, usize)> {
+    let seq: RArray = obj.try_convert()?;
+    let len = seq.len();
+    Ok((seq, len))
+}
+
 pub fn get_df(obj: Value) -> RbResult<DataFrame> {
     let rbdf = obj.funcall::<_, _, &RbDataFrame>("_df", ())?;
     Ok(rbdf.df.borrow().clone())
 }
 
-impl Into<Value> for Wrap<AnyValue<'_>> {
-    fn into(self) -> Value {
-        match self.0 {
+pub fn get_series(obj: Value) -> RbResult<Series> {
+    let rbs = obj.funcall::<_, _, &RbSeries>("_s", ())?;
+    Ok(rbs.series.borrow().clone())
+}
+
+impl TryConvert for Wrap<Utf8Chunked> {
+    fn try_convert(obj: Value) -> RbResult<Self> {
+        let (seq, len) = get_rbseq(obj)?;
+        let mut builder = Utf8ChunkedBuilder::new("", len, len * 25);
+
+        for res in seq.each() {
+            let item = res?;
+            match item.try_convert::<String>() {
+                Ok(val) => builder.append_value(&val),
+                Err(_) => builder.append_null(),
+            }
+        }
+        Ok(Wrap(builder.finish()))
+    }
+}
+
+impl TryConvert for Wrap<NullValues> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        if let Ok(s) = ob.try_convert::<String>() {
+            Ok(Wrap(NullValues::AllColumnsSingle(s)))
+        } else if let Ok(s) = ob.try_convert::<Vec<String>>() {
+            Ok(Wrap(NullValues::AllColumns(s)))
+        } else if let Ok(s) = ob.try_convert::<Vec<(String, String)>>() {
+            Ok(Wrap(NullValues::Named(s)))
+        } else {
+            Err(RbPolarsErr::other(
+                "could not extract value from null_values argument".into(),
+            ))
+        }
+    }
+}
+
+impl From<Wrap<AnyValue<'_>>> for Value {
+    fn from(w: Wrap<AnyValue<'_>>) -> Self {
+        match w.0 {
             AnyValue::UInt8(v) => Value::from(v),
             AnyValue::UInt16(v) => Value::from(v),
             AnyValue::UInt32(v) => Value::from(v),
@@ -41,6 +84,12 @@ impl Into<Value> for Wrap<AnyValue<'_>> {
     }
 }
 
+impl From<Wrap<DataType>> for Value {
+    fn from(w: Wrap<DataType>) -> Self {
+        Symbol::from(w.0.to_string()).into()
+    }
+}
+
 impl TryConvert for Wrap<DataType> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let dtype = match ob.try_convert::<String>()?.as_str() {
@@ -118,6 +167,39 @@ impl TryConvert for Wrap<ClosedWindow> {
     }
 }
 
+impl TryConvert for Wrap<CsvEncoding> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "utf8" => CsvEncoding::Utf8,
+            "utf8-lossy" => CsvEncoding::LossyUtf8,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "encoding must be one of {{'utf8', 'utf8-lossy'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<Option<IpcCompression>> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "uncompressed" => None,
+            "lz4" => Some(IpcCompression::LZ4),
+            "zstd" => Some(IpcCompression::ZSTD),
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "compression must be one of {{'uncompressed', 'lz4', 'zstd'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<JoinType> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match ob.try_convert::<String>()?.as_str() {
@@ -171,6 +253,24 @@ impl TryConvert for Wrap<NullStrategy> {
     }
 }
 
+impl TryConvert for Wrap<ParallelStrategy> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "auto" => ParallelStrategy::Auto,
+            "columns" => ParallelStrategy::Columns,
+            "row_groups" => ParallelStrategy::RowGroups,
+            "none" => ParallelStrategy::None,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "parallel must be one of {{'auto', 'columns', 'row_groups', 'none'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<QuantileInterpolOptions> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match ob.try_convert::<String>()?.as_str() {

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;
@@ -43,10 +44,94 @@ impl RbDataFrame {
         self.df.borrow().estimated_size()
     }
 
-    pub fn read_csv(rb_f: Value, has_header: bool) -> RbResult<Self> {
+    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())
@@ -61,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)?;
@@ -185,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,
@@ -240,11 +368,11 @@ impl RbDataFrame {
         Ok(())
     }
 
-    pub fn dtypes(&self) -> Vec<String> {
+    pub fn dtypes(&self) -> Vec<Value> {
         self.df
             .borrow()
             .iter()
-            .map(|s| s.dtype().to_string())
+            .map(|s| Wrap(s.dtype().clone()).into())
             .collect()
     }
 
@@ -418,7 +546,7 @@ impl RbDataFrame {
             self.df.borrow().partition_by(groups)
         }
         .map_err(RbPolarsErr::from)?;
-        Ok(out.into_iter().map(|v| RbDataFrame::new(v)).collect())
+        Ok(out.into_iter().map(RbDataFrame::new).collect())
     }
 
     pub fn shift(&self, periods: i64) -> Self {

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!();
+}