polars-df 0.23.0 → 0.24.0

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

@@ -15,7 +15,7 @@ mod serde;
 mod string;
 mod r#struct;
 
-use magnus::{RArray, prelude::*};
+use magnus::{RArray, Ruby, prelude::*};
 use polars::lazy::dsl::Expr;
 
 use crate::RbResult;
@@ -32,10 +32,26 @@ impl From<Expr> for RbExpr {
     }
 }
 
-pub fn rb_exprs_to_exprs(rb_exprs: RArray) -> RbResult<Vec<Expr>> {
-    let mut exprs = Vec::new();
-    for item in rb_exprs.into_iter() {
-        exprs.push(<&RbExpr>::try_convert(item)?.inner.clone());
+pub(crate) trait ToExprs {
+    fn to_exprs(self) -> RbResult<Vec<Expr>>;
+}
+
+impl ToExprs for RArray {
+    fn to_exprs(self) -> RbResult<Vec<Expr>> {
+        let mut exprs = Vec::new();
+        for item in self.into_iter() {
+            exprs.push(<&RbExpr>::try_convert(item)?.inner.clone());
+        }
+        Ok(exprs)
+    }
+}
+
+pub(crate) trait ToRbExprs {
+    fn to_rbexprs(self, rb: &Ruby) -> RArray;
+}
+
+impl ToRbExprs for Vec<Expr> {
+    fn to_rbexprs(self, rb: &Ruby) -> RArray {
+        rb.ary_from_iter(self.iter().map(|e| RbExpr::from(e.clone())))
     }
-    Ok(exprs)
 }

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

@@ -3,6 +3,7 @@ use polars::prelude::*;
 use polars_utils::format_pl_smallstr;
 
 use crate::RbExpr;
+use crate::utils::RubyAttach;
 
 impl RbExpr {
     pub fn name_keep(&self) -> Self {
@@ -12,14 +13,16 @@ impl RbExpr {
     pub fn name_map(&self, lambda: Proc) -> Self {
         let lambda = Opaque::from(lambda);
         let func = PlanCallback::new(move |name: PlSmallStr| {
-            let lambda = Ruby::get().unwrap().get_inner(lambda);
-            let out = lambda.call::<_, String>((name.as_str(),));
-            match out {
-                Ok(out) => Ok(format_pl_smallstr!("{}", out)),
-                Err(e) => Err(PolarsError::ComputeError(
-                    format!("Ruby function in 'name.map' produced an error: {e}.").into(),
-                )),
-            }
+            Ruby::attach(|rb| {
+                let lambda = rb.get_inner(lambda);
+                let out = lambda.call::<_, String>((name.as_str(),));
+                match out {
+                    Ok(out) => Ok(format_pl_smallstr!("{}", out)),
+                    Err(e) => Err(PolarsError::ComputeError(
+                        format!("Ruby function in 'name.map' produced an error: {e}.").into(),
+                    )),
+                }
+            })
         });
         self.inner.clone().name().map(func).into()
     }
@@ -40,6 +43,14 @@ impl RbExpr {
         self.inner.clone().name().to_uppercase().into()
     }
 
+    pub fn name_replace(&self, pattern: String, value: String, literal: bool) -> Self {
+        self.inner
+            .clone()
+            .name()
+            .replace(&pattern, &value, literal)
+            .into()
+    }
+
     pub fn name_prefix_fields(&self, prefix: String) -> Self {
         self.inner.clone().name().prefix_fields(&prefix).into()
     }

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

@@ -1,7 +1,7 @@
 use polars::prelude::*;
 
-use crate::RbExpr;
 use crate::conversion::Wrap;
+use crate::{RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
     pub fn rolling_sum(
@@ -319,6 +319,55 @@ impl RbExpr {
             .into()
     }
 
+    pub fn rolling_rank(
+        &self,
+        window_size: usize,
+        method: Wrap<RollingRankMethod>,
+        seed: Option<u64>,
+        min_samples: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_samples = min_samples.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            window_size,
+            min_periods: min_samples,
+            weights: None,
+            center,
+            fn_params: Some(RollingFnParams::Rank {
+                method: method.0,
+                seed,
+            }),
+        };
+
+        self.inner.clone().rolling_rank(options).into()
+    }
+
+    pub fn rolling_rank_by(
+        &self,
+        by: &RbExpr,
+        window_size: String,
+        method: Wrap<RollingRankMethod>,
+        seed: Option<u64>,
+        min_samples: usize,
+        closed: Wrap<ClosedWindow>,
+    ) -> RbResult<Self> {
+        let options = RollingOptionsDynamicWindow {
+            window_size: Duration::try_parse(&window_size).map_err(RbPolarsErr::from)?,
+            min_periods: min_samples,
+            closed_window: closed.0,
+            fn_params: Some(RollingFnParams::Rank {
+                method: method.0,
+                seed,
+            }),
+        };
+
+        Ok(self
+            .inner
+            .clone()
+            .rolling_rank_by(by.inner.clone(), options)
+            .into())
+    }
+
     pub fn rolling_skew(
         &self,
         window_size: usize,

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

@@ -30,7 +30,6 @@ impl RbExpr {
         self.inner.clone().str().to_date(options).into()
     }
 
-    #[allow(clippy::too_many_arguments)]
     pub fn str_to_datetime(
         &self,
         format: Option<String>,

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

@@ -1,6 +1,7 @@
 use magnus::RArray;
 
-use crate::{RbExpr, RbResult, rb_exprs_to_exprs};
+use crate::expr::ToExprs;
+use crate::{RbExpr, RbResult};
 
 impl RbExpr {
     pub fn struct_field_by_index(&self, index: i64) -> Self {
@@ -11,6 +12,10 @@ impl RbExpr {
         self.inner.clone().struct_().field_by_name(&name).into()
     }
 
+    pub fn struct_multiple_fields(&self, names: Vec<String>) -> Self {
+        self.inner.clone().struct_().field_by_names(&names).into()
+    }
+
     pub fn struct_rename_fields(&self, names: Vec<String>) -> Self {
         self.inner.clone().struct_().rename_fields(names).into()
     }
@@ -20,7 +25,7 @@ impl RbExpr {
     }
 
     pub fn struct_with_fields(&self, fields: RArray) -> RbResult<Self> {
-        let fields = rb_exprs_to_exprs(fields)?;
+        let fields = fields.to_exprs()?;
         let e = self.inner.clone().struct_().with_fields(fields);
         Ok(e.into())
     }

data/ext/polars/src/file.rs

@@ -4,69 +4,90 @@ use std::io::{Cursor, Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 
 use magnus::{Error, RString, Ruby, Value, prelude::*, value::Opaque};
-use polars::io::cloud::CloudOptions;
 use polars::io::mmap::MmapBytesReader;
 use polars::prelude::PlPath;
 use polars::prelude::file::DynWriteable;
 use polars::prelude::sync_on_close::SyncOnCloseType;
-use polars_utils::file::ClosableFile;
+use polars_utils::create_file;
+use polars_utils::file::{ClosableFile, WriteClose};
 use polars_utils::mmap::MemSlice;
 
-use crate::RbResult;
 use crate::error::RbPolarsErr;
 use crate::prelude::resolve_homedir;
+use crate::utils::RubyAttach;
+use crate::{RbErr, RbResult};
 
-#[derive(Clone)]
 pub struct RbFileLikeObject {
     inner: Opaque<Value>,
+    expects_str: bool,
+    has_flush: bool,
 }
 
+impl WriteClose for RbFileLikeObject {}
+
 impl DynWriteable for RbFileLikeObject {
     fn as_dyn_write(&self) -> &(dyn io::Write + Send + 'static) {
         self as _
     }
+
     fn as_mut_dyn_write(&mut self) -> &mut (dyn io::Write + Send + 'static) {
         self as _
     }
+
     fn close(self: Box<Self>) -> io::Result<()> {
         Ok(())
     }
+
     fn sync_on_close(&mut self, _sync_on_close: SyncOnCloseType) -> io::Result<()> {
         Ok(())
     }
 }
 
+impl Clone for RbFileLikeObject {
+    fn clone(&self) -> Self {
+        Ruby::attach(|_rb| Self {
+            // TODO clone
+            inner: self.inner,
+            expects_str: self.expects_str,
+            has_flush: self.has_flush,
+        })
+    }
+}
+
 /// 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 {
+    pub fn new(object: Value, expects_str: bool, has_flush: bool) -> Self {
         RbFileLikeObject {
             inner: object.into(),
+            expects_str,
+            has_flush,
         }
     }
 
-    pub fn as_bytes(&self) -> bytes::Bytes {
-        self.as_file_buffer().into_inner().into()
-    }
-
-    pub fn as_file_buffer(&self) -> Cursor<Vec<u8>> {
-        let bytes = Ruby::get()
-            .unwrap()
-            .get_inner(self.inner)
-            .funcall::<_, _, RString>("read", ())
-            .expect("no read method found");
-
-        let buf = unsafe { bytes.as_slice() }.to_vec();
+    pub(crate) fn to_memslice(&self) -> MemSlice {
+        Ruby::attach(|rb| {
+            let bytes = rb
+                .get_inner(self.inner)
+                .funcall::<_, _, RString>("read", ())
+                .expect("no read method found");
 
-        Cursor::new(buf)
+            let b = unsafe { bytes.as_slice() }.to_vec();
+            MemSlice::from_vec(b)
+        })
     }
 
     /// 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> {
+    pub fn ensure_requirements(
+        object: Value,
+        read: bool,
+        write: bool,
+        seek: bool,
+    ) -> RbResult<Self> {
         let ruby = Ruby::get_with(object);
 
         if read && !object.respond_to("read", false)? {
@@ -90,48 +111,61 @@ impl RbFileLikeObject {
             ));
         }
 
-        Ok(RbFileLikeObject::new(object))
+        // TODO fix
+        let expects_str = false;
+        let has_flush = object.respond_to("write", false)?;
+        Ok(RbFileLikeObject::new(object, expects_str, has_flush))
     }
 }
 
 /// 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::other(e.to_string())
+fn rberr_to_io_err(e: RbErr) -> io::Error {
+    Ruby::attach(|_rb| io::Error::other(e.to_string()))
 }
 
 impl Read for RbFileLikeObject {
     fn read(&mut self, mut buf: &mut [u8]) -> Result<usize, io::Error> {
-        let bytes = Ruby::get()
-            .unwrap()
-            .get_inner(self.inner)
-            .funcall::<_, _, RString>("read", (buf.len(),))
-            .map_err(rberr_to_io_err)?;
+        Ruby::attach(|rb| {
+            let bytes = rb
+                .get_inner(self.inner)
+                .funcall::<_, _, RString>("read", (buf.len(),))
+                .map_err(rberr_to_io_err)?;
 
-        buf.write_all(unsafe { bytes.as_slice() })?;
+            buf.write_all(unsafe { bytes.as_slice() })?;
 
-        Ok(bytes.len())
+            Ok(bytes.len())
+        })
     }
 }
 
 impl Write for RbFileLikeObject {
     fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
-        let ruby = Ruby::get().unwrap();
-        let rbbytes = ruby.str_from_slice(buf);
+        let expects_str = self.expects_str;
 
-        let number_bytes_written = ruby
-            .get_inner(self.inner)
-            .funcall::<_, _, usize>("write", (rbbytes,))
-            .map_err(rberr_to_io_err)?;
+        Ruby::attach(|rb| {
+            if expects_str {
+                todo!();
+            }
+
+            let rbbytes = rb.str_from_slice(buf);
+
+            let number_bytes_written = rb
+                .get_inner(self.inner)
+                .funcall::<_, _, usize>("write", (rbbytes,))
+                .map_err(rberr_to_io_err)?;
 
-        Ok(number_bytes_written)
+            Ok(number_bytes_written)
+        })
     }
 
     fn flush(&mut self) -> Result<(), io::Error> {
-        Ruby::get()
-            .unwrap()
-            .get_inner(self.inner)
-            .funcall::<_, _, Value>("flush", ())
-            .map_err(rberr_to_io_err)?;
+        if self.has_flush {
+            Ruby::attach(|rb| {
+                rb.get_inner(self.inner)
+                    .funcall::<_, _, Value>("flush", ())
+                    .map_err(rberr_to_io_err)
+            })?;
+        }
 
         Ok(())
     }
@@ -139,19 +173,21 @@ impl Write for RbFileLikeObject {
 
 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),
-        };
+        Ruby::attach(|rb| {
+            let (whence, offset) = match pos {
+                SeekFrom::Start(i) => (0, i as i64),
+                SeekFrom::Current(i) => (1, i),
+                SeekFrom::End(i) => (2, i),
+            };
 
-        let inner = Ruby::get().unwrap().get_inner(self.inner);
+            let inner = rb.get_inner(self.inner);
 
-        inner
-            .funcall::<_, _, Value>("seek", (offset, whence))
-            .map_err(rberr_to_io_err)?;
+            inner
+                .funcall::<_, _, Value>("seek", (offset, whence))
+                .map_err(rberr_to_io_err)?;
 
-        inner.funcall("tell", ()).map_err(rberr_to_io_err)
+            inner.funcall("tell", ()).map_err(rberr_to_io_err)
+        })
     }
 }
 
@@ -174,17 +210,18 @@ impl EitherRustRubyFile {
         }
     }
 
-    pub(crate) fn into_writeable(self) -> Box<dyn DynWriteable> {
+    #[allow(dead_code)]
+    fn into_scan_source_input(self) -> RubyScanSourceInput {
         match self {
-            Self::Rb(f) => Box::new(f),
-            Self::Rust(f) => Box::new(f),
+            EitherRustRubyFile::Rb(f) => RubyScanSourceInput::Buffer(f.to_memslice()),
+            EitherRustRubyFile::Rust(f) => RubyScanSourceInput::File(f),
         }
     }
 
-    pub fn into_dyn_writeable(self) -> Box<dyn Write> {
+    pub(crate) fn into_writeable(self) -> Box<dyn DynWriteable> {
         match self {
-            EitherRustRubyFile::Rb(f) => Box::new(f),
-            EitherRustRubyFile::Rust(f) => Box::new(f),
+            Self::Rb(f) => Box::new(f),
+            Self::Rust(f) => Box::new(f),
         }
     }
 }
@@ -193,52 +230,61 @@ pub enum RubyScanSourceInput {
     Buffer(MemSlice),
     Path(PlPath),
     #[allow(dead_code)]
-    File(File),
+    File(ClosableFile),
 }
 
 pub(crate) fn try_get_rbfile(
+    _rb: &Ruby,
     rb_f: Value,
     write: bool,
 ) -> RbResult<(EitherRustRubyFile, Option<PathBuf>)> {
-    let f = RbFileLikeObject::with_requirements(rb_f, !write, write, !write)?;
+    let f = RbFileLikeObject::ensure_requirements(rb_f, !write, write, !write)?;
     Ok((EitherRustRubyFile::Rb(f), None))
 }
 
 pub fn get_ruby_scan_source_input(rb_f: Value, write: bool) -> RbResult<RubyScanSourceInput> {
-    if let Ok(s) = String::try_convert(rb_f) {
-        let mut file_path = PlPath::new(&s);
-        if let Some(p) = file_path.as_ref().as_local_path() {
-            if p.starts_with("~/") {
+    Ruby::attach(|_rb| {
+        if let Ok(s) = String::try_convert(rb_f) {
+            let mut file_path = PlPath::new(&s);
+            if let Some(p) = file_path.as_ref().as_local_path()
+                && p.starts_with("~/")
+            {
                 file_path = PlPath::Local(resolve_homedir(&p).into());
             }
+            Ok(RubyScanSourceInput::Path(file_path))
+        } else {
+            let f = RbFileLikeObject::ensure_requirements(rb_f, !write, write, !write)?;
+            Ok(RubyScanSourceInput::Buffer(f.to_memslice()))
         }
-        Ok(RubyScanSourceInput::Path(file_path))
-    } else {
-        let f = RbFileLikeObject::with_requirements(rb_f, !write, write, !write)?;
-        Ok(RubyScanSourceInput::Buffer(MemSlice::from_bytes(
-            f.as_bytes(),
-        )))
-    }
+    })
+}
+
+pub fn get_either_buffer_or_path(
+    rb_f: Value,
+    write: bool,
+) -> RbResult<(EitherRustRubyFile, Option<PathBuf>)> {
+    Ruby::attach(|rb| {
+        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 write {
+                create_file(&file_path).map_err(RbPolarsErr::from)?
+            } else {
+                polars_utils::open_file(&file_path).map_err(RbPolarsErr::from)?
+            };
+            Ok((EitherRustRubyFile::Rust(f.into()), Some(file_path)))
+        } else {
+            try_get_rbfile(rb, rb_f, write)
+        }
+    })
 }
 
 ///
 /// # 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::from)?
-        } else {
-            polars_utils::open_file(&file_path).map_err(RbPolarsErr::from)?
-        };
-        Ok(EitherRustRubyFile::Rust(f.into()))
-    } else {
-        let f = RbFileLikeObject::with_requirements(rb_f, !truncate, truncate, !truncate)?;
-        Ok(EitherRustRubyFile::Rb(f))
-    }
+/// * `write` - open for writing; will truncate existing file and create new file if not.
+pub fn get_either_file(rb_f: Value, write: bool) -> RbResult<EitherRustRubyFile> {
+    Ok(get_either_buffer_or_path(rb_f, write)?.0)
 }
 
 pub fn get_file_like(f: Value, truncate: bool) -> RbResult<Box<dyn FileLike>> {
@@ -265,22 +311,9 @@ pub fn get_mmap_bytes_reader_and_path<'a>(
 ) -> RbResult<(Box<dyn MmapBytesReader + 'a>, Option<PathBuf>)> {
     match rb_f {
         RbReadBytes::Bytes(v) => Ok((Box::new(Cursor::new(unsafe { v.as_slice() })), None)),
-        RbReadBytes::Other(v) => {
-            let path = PathBuf::try_convert(*v)?;
-            let f = File::open(&path).map_err(|e| {
-                Error::new(
-                    Ruby::get().unwrap().exception_runtime_error(),
-                    e.to_string(),
-                )
-            })?;
-            Ok((Box::new(f), Some(path)))
-        }
+        RbReadBytes::Other(v) => match get_either_buffer_or_path(*v, false)? {
+            (EitherRustRubyFile::Rust(f), path) => Ok((Box::new(f), path)),
+            (EitherRustRubyFile::Rb(f), path) => Ok((Box::new(Cursor::new(f.to_memslice())), path)),
+        },
     }
 }
-
-pub fn try_get_writeable(
-    rb_f: Value,
-    _cloud_options: Option<&CloudOptions>,
-) -> RbResult<Box<dyn Write>> {
-    Ok(get_either_file(rb_f, true)?.into_dyn_writeable())
-}

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

@@ -1,41 +1,42 @@
 use magnus::RArray;
 use polars::lazy::dsl;
 
-use crate::rb_exprs_to_exprs;
-use crate::{RbExpr, RbPolarsErr, RbResult};
+use crate::error::RbPolarsErr;
+use crate::expr::ToExprs;
+use crate::{RbExpr, RbResult};
 
 pub fn all_horizontal(exprs: RArray) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
     let e = dsl::all_horizontal(exprs).map_err(RbPolarsErr::from)?;
     Ok(e.into())
 }
 
 pub fn any_horizontal(exprs: RArray) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
     let e = dsl::any_horizontal(exprs).map_err(RbPolarsErr::from)?;
     Ok(e.into())
 }
 
 pub fn max_horizontal(exprs: RArray) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
     let e = dsl::max_horizontal(exprs).map_err(RbPolarsErr::from)?;
     Ok(e.into())
 }
 
 pub fn min_horizontal(exprs: RArray) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
     let e = dsl::min_horizontal(exprs).map_err(RbPolarsErr::from)?;
     Ok(e.into())
 }
 
 pub fn sum_horizontal(exprs: RArray, ignore_nulls: bool) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
     let e = dsl::sum_horizontal(exprs, ignore_nulls).map_err(RbPolarsErr::from)?;
     Ok(e.into())
 }
 
 pub fn mean_horizontal(exprs: RArray, ignore_nulls: bool) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
     let e = dsl::mean_horizontal(exprs, ignore_nulls).map_err(RbPolarsErr::from)?;
     Ok(e.into())
 }