polars-df 0.8.0 → 0.10.0

data/ext/polars/src/file.rs

@@ -1,20 +1,167 @@
-use magnus::{exception, prelude::*, Error, RString, Value};
-use polars::io::mmap::MmapBytesReader;
 use std::fs::File;
-use std::io::Cursor;
+use std::io;
+use std::io::{BufReader, Cursor, Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 
+use magnus::{exception, prelude::*, Error, RString, Value};
+use polars::io::mmap::MmapBytesReader;
+
+use crate::error::RbPolarsErr;
+use crate::prelude::resolve_homedir;
 use crate::RbResult;
 
-pub fn get_file_like(f: Value, truncate: bool) -> RbResult<File> {
-    let str_slice = PathBuf::try_convert(f)?;
-    let f = if truncate {
-        File::create(str_slice)
-            .map_err(|e| Error::new(exception::runtime_error(), e.to_string()))?
+#[derive(Clone)]
+pub struct RbFileLikeObject {
+    inner: Value,
+}
+
+/// Wraps a `Value`, and implements read, seek, and write for it.
+impl RbFileLikeObject {
+    /// Creates an instance of a `RbFileLikeObject` from a `Value`.
+    /// To assert the object has the required methods methods,
+    /// instantiate it with `RbFileLikeObject::require`
+    pub fn new(object: Value) -> Self {
+        RbFileLikeObject { inner: object }
+    }
+
+    pub fn as_buffer(&self) -> std::io::Cursor<Vec<u8>> {
+        let data = self.as_file_buffer().into_inner();
+        std::io::Cursor::new(data)
+    }
+
+    pub fn as_file_buffer(&self) -> Cursor<Vec<u8>> {
+        let bytes = self
+            .inner
+            .funcall::<_, _, RString>("read", ())
+            .expect("no read method found");
+
+        let buf = unsafe { bytes.as_slice() }.to_vec();
+
+        Cursor::new(buf)
+    }
+
+    /// Same as `RbFileLikeObject::new`, but validates that the underlying
+    /// ruby object has a `read`, `write`, and `seek` methods in respect to parameters.
+    /// Will return a `TypeError` if object does not have `read`, `seek`, and `write` methods.
+    pub fn with_requirements(object: Value, read: bool, write: bool, seek: bool) -> RbResult<Self> {
+        if read && !object.respond_to("read", false)? {
+            return Err(Error::new(
+                exception::type_error(),
+                "Object does not have a .read() method.",
+            ));
+        }
+
+        if seek && !object.respond_to("seek", false)? {
+            return Err(Error::new(
+                exception::type_error(),
+                "Object does not have a .seek() method.",
+            ));
+        }
+
+        if write && !object.respond_to("write", false)? {
+            return Err(Error::new(
+                exception::type_error(),
+                "Object does not have a .write() method.",
+            ));
+        }
+
+        Ok(RbFileLikeObject::new(object))
+    }
+}
+
+/// Extracts a string repr from, and returns an IO error to send back to rust.
+fn rberr_to_io_err(e: Error) -> io::Error {
+    io::Error::new(io::ErrorKind::Other, e.to_string())
+}
+
+impl Read for RbFileLikeObject {
+    fn read(&mut self, mut buf: &mut [u8]) -> Result<usize, io::Error> {
+        let bytes = self
+            .inner
+            .funcall::<_, _, RString>("read", (buf.len(),))
+            .map_err(rberr_to_io_err)?;
+
+        buf.write_all(unsafe { bytes.as_slice() })?;
+
+        Ok(bytes.len())
+    }
+}
+
+impl Write for RbFileLikeObject {
+    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
+        let rbbytes = RString::from_slice(buf);
+
+        let number_bytes_written = self
+            .inner
+            .funcall::<_, _, usize>("write", (rbbytes,))
+            .map_err(rberr_to_io_err)?;
+
+        Ok(number_bytes_written)
+    }
+
+    fn flush(&mut self) -> Result<(), io::Error> {
+        self.inner
+            .funcall::<_, _, Value>("flush", ())
+            .map_err(rberr_to_io_err)?;
+
+        Ok(())
+    }
+}
+
+impl Seek for RbFileLikeObject {
+    fn seek(&mut self, pos: SeekFrom) -> Result<u64, io::Error> {
+        let (whence, offset) = match pos {
+            SeekFrom::Start(i) => (0, i as i64),
+            SeekFrom::Current(i) => (1, i),
+            SeekFrom::End(i) => (2, i),
+        };
+
+        let new_position = self
+            .inner
+            .funcall("seek", (offset, whence))
+            .map_err(rberr_to_io_err)?;
+
+        Ok(new_position)
+    }
+}
+
+pub trait FileLike: Read + Write + Seek {}
+
+impl FileLike for File {}
+impl FileLike for RbFileLikeObject {}
+
+pub enum EitherRustRubyFile {
+    Rb(RbFileLikeObject),
+    Rust(BufReader<File>),
+}
+
+///
+/// # Arguments
+/// * `truncate` - open or create a new file.
+pub fn get_either_file(rb_f: Value, truncate: bool) -> RbResult<EitherRustRubyFile> {
+    if let Ok(rstring) = RString::try_convert(rb_f) {
+        let s = unsafe { rstring.as_str() }?;
+        let file_path = std::path::Path::new(&s);
+        let file_path = resolve_homedir(file_path);
+        let f = if truncate {
+            File::create(file_path).map_err(RbPolarsErr::io)?
+        } else {
+            polars_utils::open_file(&file_path).map_err(RbPolarsErr::from)?
+        };
+        let reader = BufReader::new(f);
+        Ok(EitherRustRubyFile::Rust(reader))
     } else {
-        File::open(str_slice).map_err(|e| Error::new(exception::runtime_error(), e.to_string()))?
-    };
-    Ok(f)
+        let f = RbFileLikeObject::with_requirements(rb_f, !truncate, truncate, !truncate)?;
+        Ok(EitherRustRubyFile::Rb(f))
+    }
+}
+
+pub fn get_file_like(f: Value, truncate: bool) -> RbResult<Box<dyn FileLike>> {
+    use EitherRustRubyFile::*;
+    match get_either_file(f, truncate)? {
+        Rb(f) => Ok(Box::new(f)),
+        Rust(f) => Ok(Box::new(f.into_inner())),
+    }
 }
 
 pub fn get_mmap_bytes_reader(rb_f: Value) -> RbResult<Box<dyn MmapBytesReader>> {

data/ext/polars/src/functions/aggregation.rs

@@ -33,3 +33,9 @@ pub fn sum_horizontal(exprs: RArray) -> RbResult<RbExpr> {
     let e = dsl::sum_horizontal(exprs).map_err(RbPolarsErr::from)?;
     Ok(e.into())
 }
+
+pub fn mean_horizontal(exprs: RArray) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let e = dsl::mean_horizontal(exprs).map_err(RbPolarsErr::from)?;
+    Ok(e.into())
+}

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

@@ -17,9 +17,62 @@ macro_rules! set_unwrapped_or_0 {
     };
 }
 
-pub fn arg_sort_by(by: RArray, descending: Vec<bool>) -> RbResult<RbExpr> {
+pub fn rolling_corr(
+    x: &RbExpr,
+    y: &RbExpr,
+    window_size: IdxSize,
+    min_periods: IdxSize,
+    ddof: u8,
+) -> RbExpr {
+    dsl::rolling_corr(
+        x.inner.clone(),
+        y.inner.clone(),
+        RollingCovOptions {
+            min_periods,
+            window_size,
+            ddof,
+        },
+    )
+    .into()
+}
+
+pub fn rolling_cov(
+    x: &RbExpr,
+    y: &RbExpr,
+    window_size: IdxSize,
+    min_periods: IdxSize,
+    ddof: u8,
+) -> RbExpr {
+    dsl::rolling_cov(
+        x.inner.clone(),
+        y.inner.clone(),
+        RollingCovOptions {
+            min_periods,
+            window_size,
+            ddof,
+        },
+    )
+    .into()
+}
+
+pub fn arg_sort_by(
+    by: RArray,
+    descending: Vec<bool>,
+    nulls_last: bool,
+    multithreaded: bool,
+    maintain_order: bool,
+) -> RbResult<RbExpr> {
     let by = rb_exprs_to_exprs(by)?;
-    Ok(dsl::arg_sort_by(by, &descending).into())
+    Ok(dsl::arg_sort_by(
+        by,
+        SortMultipleOptions {
+            descending,
+            nulls_last,
+            multithreaded,
+            maintain_order,
+        },
+    )
+    .into())
 }
 
 pub fn arg_where(condition: &RbExpr) -> RbExpr {
@@ -83,6 +136,47 @@ pub fn concat_lf(
     Ok(lf.into())
 }
 
+pub fn concat_list(s: RArray) -> RbResult<RbExpr> {
+    let s = rb_exprs_to_exprs(s)?;
+    let expr = dsl::concat_list(s).map_err(RbPolarsErr::from)?;
+    Ok(expr.into())
+}
+
+pub fn concat_str(s: RArray, separator: String, ignore_nulls: bool) -> RbResult<RbExpr> {
+    let s = rb_exprs_to_exprs(s)?;
+    Ok(dsl::concat_str(s, &separator, ignore_nulls).into())
+}
+
+pub fn len() -> RbExpr {
+    dsl::len().into()
+}
+
+pub fn cov(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
+    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone(), ddof).into()
+}
+
+pub fn arctan2(y: &RbExpr, x: &RbExpr) -> RbExpr {
+    y.inner.clone().arctan2(x.inner.clone()).into()
+}
+
+pub fn arctan2d(y: &RbExpr, x: &RbExpr) -> RbExpr {
+    y.inner.clone().arctan2(x.inner.clone()).degrees().into()
+}
+
+pub fn cum_fold(
+    acc: &RbExpr,
+    lambda: Value,
+    exprs: RArray,
+    include_init: bool,
+) -> RbResult<RbExpr> {
+    let exprs = rb_exprs_to_exprs(exprs)?;
+    let lambda = Opaque::from(lambda);
+
+    let func =
+        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
+    Ok(dsl::cum_fold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+}
+
 pub fn concat_lf_diagonal(
     lfs: RArray,
     rechunk: bool,
@@ -110,6 +204,19 @@ pub fn concat_lf_diagonal(
     Ok(lf.into())
 }
 
+pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
+    dsl::dtype_cols(dtypes).into()
+}
+
+pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
+    let dtypes = dtypes
+        .each()
+        .map(|v| Wrap::<DataType>::try_convert(v?))
+        .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
+    let dtypes = vec_extract_wrapped(dtypes);
+    Ok(crate::functions::lazy::dtype_cols(dtypes))
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn duration(
     weeks: Option<&RbExpr>,
@@ -146,38 +253,21 @@ pub fn duration(
     dsl::duration(args).into()
 }
 
-pub fn count() -> RbExpr {
-    dsl::count().into()
-}
-
 pub fn first() -> RbExpr {
     dsl::first().into()
 }
 
-pub fn last() -> RbExpr {
-    dsl::last().into()
-}
-
-pub fn dtype_cols(dtypes: Vec<DataType>) -> RbExpr {
-    dsl::dtype_cols(dtypes).into()
-}
-
 pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
     let lambda = Opaque::from(lambda);
 
     let func =
         move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
-    Ok(polars::lazy::dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
+    Ok(dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
 }
 
-pub fn cumfold(acc: &RbExpr, lambda: Value, exprs: RArray, include_init: bool) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
-    let lambda = Opaque::from(lambda);
-
-    let func =
-        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
-    Ok(polars::lazy::dsl::cum_fold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+pub fn last() -> RbExpr {
+    dsl::last().into()
 }
 
 pub fn lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
@@ -219,6 +309,10 @@ pub fn lit(value: Value, allow_object: bool) -> RbResult<RbExpr> {
     }
 }
 
+pub fn pearson_corr(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
+    dsl::pearson_corr(a.inner.clone(), b.inner.clone(), ddof).into()
+}
+
 pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbResult<RbExpr> {
     let mut value = value.inner.clone();
     let n = n.inner.clone();
@@ -228,7 +322,7 @@ pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbRe
     }
 
     if let Expr::Literal(lv) = &value {
-        let av = lv.to_anyvalue().unwrap();
+        let av = lv.to_any_value().unwrap();
         // Integer inputs that fit in Int32 are parsed as such
         if let DataType::Int64 = av.dtype() {
             let int_value = av.try_extract::<i64>().unwrap();
@@ -240,35 +334,11 @@ pub fn repeat(value: &RbExpr, n: &RbExpr, dtype: Option<Wrap<DataType>>) -> RbRe
     Ok(dsl::repeat(value, n).into())
 }
 
-pub fn pearson_corr(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
-    polars::lazy::dsl::pearson_corr(a.inner.clone(), b.inner.clone(), ddof).into()
-}
-
 pub fn spearman_rank_corr(a: &RbExpr, b: &RbExpr, ddof: u8, propagate_nans: bool) -> RbExpr {
-    polars::lazy::dsl::spearman_rank_corr(a.inner.clone(), b.inner.clone(), ddof, propagate_nans)
-        .into()
-}
-
-pub fn cov(a: &RbExpr, b: &RbExpr, ddof: u8) -> RbExpr {
-    polars::lazy::dsl::cov(a.inner.clone(), b.inner.clone(), ddof).into()
-}
-
-pub fn concat_str(s: RArray, sep: String) -> RbResult<RbExpr> {
-    let s = rb_exprs_to_exprs(s)?;
-    Ok(dsl::concat_str(s, &sep).into())
+    dsl::spearman_rank_corr(a.inner.clone(), b.inner.clone(), ddof, propagate_nans).into()
 }
 
-pub fn concat_lst(s: RArray) -> RbResult<RbExpr> {
-    let s = rb_exprs_to_exprs(s)?;
-    let expr = dsl::concat_list(s).map_err(RbPolarsErr::from)?;
+pub fn sql_expr(sql: String) -> RbResult<RbExpr> {
+    let expr = polars::sql::sql_expr(sql).map_err(RbPolarsErr::from)?;
     Ok(expr.into())
 }
-
-pub fn dtype_cols2(dtypes: RArray) -> RbResult<RbExpr> {
-    let dtypes = dtypes
-        .each()
-        .map(|v| Wrap::<DataType>::try_convert(v?))
-        .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
-    let dtypes = vec_extract_wrapped(dtypes);
-    Ok(crate::functions::lazy::dtype_cols(dtypes))
-}

data/ext/polars/src/functions/meta.rs

@@ -7,7 +7,7 @@ use polars_core::POOL;
 use crate::conversion::Wrap;
 use crate::{RbResult, RbValueError};
 
-pub fn get_idx_type() -> Value {
+pub fn get_index_type() -> Value {
     Wrap(IDX_DTYPE).into_value()
 }
 
@@ -36,3 +36,47 @@ pub fn get_float_fmt() -> RbResult<String> {
     };
     Ok(strfmt.to_string())
 }
+
+pub fn set_float_precision(precision: Option<usize>) -> RbResult<()> {
+    use polars_core::fmt::set_float_precision;
+    set_float_precision(precision);
+    Ok(())
+}
+
+pub fn get_float_precision() -> RbResult<Option<usize>> {
+    use polars_core::fmt::get_float_precision;
+    Ok(get_float_precision())
+}
+
+pub fn set_thousands_separator(sep: Option<char>) -> RbResult<()> {
+    use polars_core::fmt::set_thousands_separator;
+    set_thousands_separator(sep);
+    Ok(())
+}
+
+pub fn get_thousands_separator() -> RbResult<Option<String>> {
+    use polars_core::fmt::get_thousands_separator;
+    Ok(Some(get_thousands_separator()))
+}
+
+pub fn set_decimal_separator(sep: Option<char>) -> RbResult<()> {
+    use polars_core::fmt::set_decimal_separator;
+    set_decimal_separator(sep);
+    Ok(())
+}
+
+pub fn get_decimal_separator() -> RbResult<Option<char>> {
+    use polars_core::fmt::get_decimal_separator;
+    Ok(Some(get_decimal_separator()))
+}
+
+pub fn set_trim_decimal_zeros(trim: Option<bool>) -> RbResult<()> {
+    use polars_core::fmt::set_trim_decimal_zeros;
+    set_trim_decimal_zeros(trim);
+    Ok(())
+}
+
+pub fn get_trim_decimal_zeros() -> RbResult<Option<bool>> {
+    use polars_core::fmt::get_trim_decimal_zeros;
+    Ok(Some(get_trim_decimal_zeros()))
+}

data/ext/polars/src/functions/string_cache.rs

@@ -1,3 +1,7 @@
+use crate::RbResult;
+use magnus::{RArray, Ruby, Value};
+use polars_core::StringCacheHolder;
+
 pub fn enable_string_cache() {
     polars_core::enable_string_cache()
 }
@@ -9,3 +13,13 @@ pub fn disable_string_cache() {
 pub fn using_string_cache() -> bool {
     polars_core::using_string_cache()
 }
+
+#[magnus::wrap(class = "Polars::RbStringCacheHolder")]
+pub struct RbStringCacheHolder {}
+
+impl RbStringCacheHolder {
+    pub fn hold() -> RbResult<Value> {
+        let _hold = StringCacheHolder::hold();
+        Ruby::get().unwrap().yield_splat(RArray::new())
+    }
+}

data/ext/polars/src/functions/whenthen.rs

@@ -2,42 +2,72 @@ use polars::lazy::dsl;
 
 use crate::RbExpr;
 
+pub fn when(condition: &RbExpr) -> RbWhen {
+    RbWhen {
+        inner: dsl::when(condition.inner.clone()),
+    }
+}
+
 #[magnus::wrap(class = "Polars::RbWhen")]
 #[derive(Clone)]
 pub struct RbWhen {
     pub inner: dsl::When,
 }
 
-impl From<dsl::When> for RbWhen {
-    fn from(inner: dsl::When) -> Self {
-        RbWhen { inner }
-    }
-}
-
-#[magnus::wrap(class = "Polars::RbWhenThen")]
+#[magnus::wrap(class = "Polars::RbThen")]
 #[derive(Clone)]
 pub struct RbThen {
     pub inner: dsl::Then,
 }
 
-impl From<dsl::Then> for RbThen {
-    fn from(inner: dsl::Then) -> Self {
-        RbThen { inner }
-    }
+#[magnus::wrap(class = "Polars::RbChainedWhen")]
+#[derive(Clone)]
+pub struct RbChainedWhen {
+    pub inner: dsl::ChainedWhen,
+}
+
+#[magnus::wrap(class = "Polars::RbChainedThen")]
+#[derive(Clone)]
+pub struct RbChainedThen {
+    pub inner: dsl::ChainedThen,
 }
 
 impl RbWhen {
-    pub fn then(&self, expr: &RbExpr) -> RbThen {
-        self.inner.clone().then(expr.inner.clone()).into()
+    pub fn then(&self, statement: &RbExpr) -> RbThen {
+        RbThen {
+            inner: self.inner.clone().then(statement.inner.clone()),
+        }
     }
 }
 
 impl RbThen {
-    pub fn overwise(&self, expr: &RbExpr) -> RbExpr {
-        self.inner.clone().otherwise(expr.inner.clone()).into()
+    pub fn when(&self, condition: &RbExpr) -> RbChainedWhen {
+        RbChainedWhen {
+            inner: self.inner.clone().when(condition.inner.clone()),
+        }
+    }
+
+    pub fn otherwise(&self, statement: &RbExpr) -> RbExpr {
+        self.inner.clone().otherwise(statement.inner.clone()).into()
     }
 }
 
-pub fn when(predicate: &RbExpr) -> RbWhen {
-    dsl::when(predicate.inner.clone()).into()
+impl RbChainedWhen {
+    pub fn then(&self, statement: &RbExpr) -> RbChainedThen {
+        RbChainedThen {
+            inner: self.inner.clone().then(statement.inner.clone()),
+        }
+    }
+}
+
+impl RbChainedThen {
+    pub fn when(&self, condition: &RbExpr) -> RbChainedWhen {
+        RbChainedWhen {
+            inner: self.inner.clone().when(condition.inner.clone()),
+        }
+    }
+
+    pub fn otherwise(&self, statement: &RbExpr) -> RbExpr {
+        self.inner.clone().otherwise(statement.inner.clone()).into()
+    }
 }