polars-df 0.9.0 → 0.11.0

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

@@ -7,187 +7,281 @@ use crate::RbExpr;
 impl RbExpr {
     pub fn rolling_sum(
         &self,
-        window_size: String,
+        window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: usize,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let options = RollingOptionsFixedWindow {
+            window_size,
             weights,
             min_periods,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             ..Default::default()
         };
         self.inner.clone().rolling_sum(options).into()
     }
 
-    pub fn rolling_min(
+    pub fn rolling_sum_by(
         &self,
+        by: &RbExpr,
         window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+        self.inner
+            .clone()
+            .rolling_sum_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_min(
+        &self,
+        window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: usize,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let options = RollingOptionsFixedWindow {
+            window_size,
             weights,
             min_periods,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             ..Default::default()
         };
         self.inner.clone().rolling_min(options).into()
     }
 
-    pub fn rolling_max(
+    pub fn rolling_min_by(
         &self,
+        by: &RbExpr,
         window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+        self.inner
+            .clone()
+            .rolling_min_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_max(
+        &self,
+        window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: usize,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let options = RollingOptionsFixedWindow {
+            window_size: window_size,
             weights,
             min_periods,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             ..Default::default()
         };
         self.inner.clone().rolling_max(options).into()
     }
 
-    pub fn rolling_mean(
+    pub fn rolling_max_by(
         &self,
+        by: &RbExpr,
         window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+        self.inner
+            .clone()
+            .rolling_max_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_mean(
+        &self,
+        window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: usize,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let options = RollingOptionsFixedWindow {
+            window_size,
             weights,
             min_periods,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             ..Default::default()
         };
 
         self.inner.clone().rolling_mean(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub fn rolling_std(
+    pub fn rolling_mean_by(
         &self,
+        by: &RbExpr,
         window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+
+        self.inner
+            .clone()
+            .rolling_mean_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_std(
+        &self,
+        window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: usize,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
         ddof: u8,
-        warn_if_unsorted: bool,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let options = RollingOptionsFixedWindow {
+            window_size,
             weights,
             min_periods,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
-            warn_if_unsorted,
         };
 
         self.inner.clone().rolling_std(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub fn rolling_var(
+    pub fn rolling_std_by(
         &self,
+        by: &RbExpr,
         window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+        ddof: u8,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+        };
+
+        self.inner
+            .clone()
+            .rolling_std_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_var(
+        &self,
+        window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: usize,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
         ddof: u8,
-        warn_if_unsorted: bool,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let options = RollingOptionsFixedWindow {
+            window_size,
             weights,
             min_periods,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
-            warn_if_unsorted,
         };
 
         self.inner.clone().rolling_var(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub fn rolling_median(
+    pub fn rolling_var_by(
         &self,
+        by: &RbExpr,
         window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+        ddof: u8,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+        };
+
+        self.inner
+            .clone()
+            .rolling_var_by(by.inner.clone(), options)
+            .into()
+    }
+
+    pub fn rolling_median(
+        &self,
+        window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: usize,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-        warn_if_unsorted: bool,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
-            weights,
+        let options = RollingOptionsFixedWindow {
+            window_size,
             min_periods,
+            weights,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             fn_params: None,
-            warn_if_unsorted,
         };
         self.inner.clone().rolling_median(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
+    pub fn rolling_median_by(
+        &self,
+        by: &RbExpr,
+        window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+        self.inner
+            .clone()
+            .rolling_median_by(by.inner.clone(), options)
+            .into()
+    }
+
     pub fn rolling_quantile(
         &self,
         quantile: f64,
         interpolation: Wrap<QuantileInterpolOptions>,
-        window_size: String,
+        window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: usize,
         center: bool,
-        by: Option<String>,
-        closed: Option<Wrap<ClosedWindow>>,
-        warn_if_unsorted: bool,
     ) -> Self {
-        let options = RollingOptions {
-            window_size: Duration::parse(&window_size),
+        let options = RollingOptionsFixedWindow {
+            window_size,
             weights,
             min_periods,
             center,
-            by,
-            closed_window: closed.map(|c| c.0),
             fn_params: None,
-            warn_if_unsorted,
         };
 
         self.inner
@@ -196,6 +290,28 @@ impl RbExpr {
             .into()
     }
 
+    pub fn rolling_quantile_by(
+        &self,
+        by: &RbExpr,
+        quantile: f64,
+        interpolation: Wrap<QuantileInterpolOptions>,
+        window_size: String,
+        min_periods: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> Self {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::parse(&window_size),
+            min_periods,
+            closed_window: closed.0,
+            fn_params: None,
+        };
+
+        self.inner
+            .clone()
+            .rolling_quantile_by(by.inner.clone(), interpolation.0, quantile, options)
+            .into()
+    }
+
     pub fn rolling_skew(&self, window_size: usize, bias: bool) -> Self {
         self.inner.clone().rolling_skew(window_size, bias).into()
     }

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

@@ -244,12 +244,12 @@ impl RbExpr {
             .into()
     }
 
-    pub fn str_to_integer(&self, base: u32, strict: bool) -> Self {
+    pub fn str_to_integer(&self, base: &Self, strict: bool) -> Self {
         self.inner
             .clone()
             .str()
-            .to_integer(base, strict)
-            .with_fmt("str.parse_int")
+            .to_integer(base.inner.clone(), strict)
+            .with_fmt("str.to_integer")
             .into()
     }
 
@@ -259,39 +259,18 @@ impl RbExpr {
         infer_schema_len: Option<usize>,
     ) -> Self {
         let dtype = dtype.map(|wrap| wrap.0);
-
-        let output_type = match dtype.clone() {
-            Some(dtype) => GetOutput::from_type(dtype),
-            None => GetOutput::from_type(DataType::Unknown),
-        };
-
-        let function = move |s: Series| {
-            let ca = s.str()?;
-            match ca.json_decode(dtype.clone(), infer_schema_len) {
-                Ok(ca) => Ok(Some(ca.into_series())),
-                Err(e) => Err(PolarsError::ComputeError(format!("{e:?}").into())),
-            }
-        };
-
-        self.clone()
-            .inner
-            .map(function, output_type)
-            .with_fmt("str.json_decode")
+        self.inner
+            .clone()
+            .str()
+            .json_decode(dtype, infer_schema_len)
             .into()
     }
 
-    pub fn str_json_path_match(&self, pat: String) -> Self {
-        let function = move |s: Series| {
-            let ca = s.str()?;
-            match ca.json_path_match(&pat) {
-                Ok(ca) => Ok(Some(ca.into_series())),
-                Err(e) => Err(PolarsError::ComputeError(format!("{:?}", e).into())),
-            }
-        };
-        self.clone()
-            .inner
-            .map(function, GetOutput::from_type(DataType::String))
-            .with_fmt("str.json_path_match")
+    pub fn str_json_path_match(&self, pat: &Self) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .json_path_match(pat.inner.clone())
             .into()
     }
 

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/lazy.rs

@@ -55,9 +55,24 @@ pub fn rolling_cov(
     .into()
 }
 
-pub fn arg_sort_by(by: RArray, descending: Vec<bool>) -> RbResult<RbExpr> {
+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 {
@@ -115,6 +130,7 @@ pub fn concat_lf(
             rechunk,
             parallel,
             to_supertypes,
+            ..Default::default()
         },
     )
     .map_err(RbPolarsErr::from)?;
@@ -183,6 +199,7 @@ pub fn concat_lf_diagonal(
             rechunk,
             parallel,
             to_supertypes,
+            ..Default::default()
         },
     )
     .map_err(RbPolarsErr::from)?;
@@ -324,6 +341,6 @@ pub fn spearman_rank_corr(a: &RbExpr, b: &RbExpr, ddof: u8, propagate_nans: bool
 }
 
 pub fn sql_expr(sql: String) -> RbResult<RbExpr> {
-    let expr = polars::sql::sql_expr(&sql).map_err(RbPolarsErr::from)?;
+    let expr = polars::sql::sql_expr(sql).map_err(RbPolarsErr::from)?;
     Ok(expr.into())
 }