polars-df 0.15.0 → 0.26.0

data/ext/polars/src/file.rs

@@ -3,16 +3,49 @@ use std::io;
 use std::io::{Cursor, Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 
-use magnus::{exception, prelude::*, Error, RString, Value};
+use magnus::{Error, RString, Ruby, Value, prelude::*, value::Opaque};
 use polars::io::mmap::MmapBytesReader;
+use polars::prelude::PlRefPath;
+use polars::prelude::file::{Writeable, WriteableTrait};
+use polars_buffer::Buffer;
+use polars_utils::create_file;
 
 use crate::error::RbPolarsErr;
 use crate::prelude::resolve_homedir;
-use crate::RbResult;
+use crate::ruby::gvl::GvlExt;
+use crate::ruby::thread::{is_non_ruby_thread, run_in_ruby_thread, start_background_ruby_thread};
+use crate::utils::to_rb_err;
+use crate::{RbErr, RbResult};
 
-#[derive(Clone)]
 pub struct RbFileLikeObject {
-    inner: Value,
+    inner: Opaque<Value>,
+    expects_str: bool,
+    has_flush: bool,
+}
+
+impl WriteableTrait for RbFileLikeObject {
+    fn close(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+
+    fn sync_all(&self) -> std::io::Result<()> {
+        self.flush()
+    }
+
+    fn sync_data(&self) -> std::io::Result<()> {
+        self.flush()
+    }
+}
+
+impl Clone for RbFileLikeObject {
+    fn clone(&self) -> Self {
+        // Skip attach
+        Self {
+            inner: self.inner,
+            expects_str: self.expects_str,
+            has_flush: self.has_flush,
+        }
+    }
 }
 
 /// Wraps a `Value`, and implements read, seek, and write for it.
@@ -20,107 +53,142 @@ 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_bytes(&self) -> bytes::Bytes {
-        self.as_file_buffer().into_inner().into()
+    pub fn new(object: Value, expects_str: bool, has_flush: bool) -> Self {
+        RbFileLikeObject {
+            inner: object.into(),
+            expects_str,
+            has_flush,
+        }
     }
 
-    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();
+    pub(crate) fn to_buffer(&self) -> Buffer<u8> {
+        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();
+            Buffer::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<()> {
+        let ruby = Ruby::get_with(object);
+
         if read && !object.respond_to("read", false)? {
             return Err(Error::new(
-                exception::type_error(),
+                ruby.exception_type_error(),
                 "Object does not have a .read() method.",
             ));
         }
 
         if seek && !object.respond_to("seek", false)? {
             return Err(Error::new(
-                exception::type_error(),
+                ruby.exception_type_error(),
                 "Object does not have a .seek() method.",
             ));
         }
 
         if write && !object.respond_to("write", false)? {
             return Err(Error::new(
-                exception::type_error(),
+                ruby.exception_type_error(),
                 "Object does not have a .write() method.",
             ));
         }
 
-        Ok(RbFileLikeObject::new(object))
+        Ok(())
+    }
+
+    fn flush(&self) -> std::io::Result<()> {
+        if self.has_flush {
+            if is_non_ruby_thread() {
+                let self2 = self.clone();
+                return run_in_ruby_thread(move || self2.flush());
+            }
+
+            Ruby::attach(|rb| {
+                rb.get_inner(self.inner)
+                    .funcall::<_, _, Value>("flush", ())
+                    .map_err(rberr_to_io_err)
+            })?;
+        }
+
+        Ok(())
     }
 }
 
 /// 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())
+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 = 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 rbbytes = RString::from_slice(buf);
+        if is_non_ruby_thread() {
+            let mut self2 = self.clone();
+            let buf2 = buf.to_vec();
+            return run_in_ruby_thread(move || self2.write(&buf2));
+        }
+
+        let expects_str = self.expects_str;
+
+        Ruby::attach(|rb| {
+            if expects_str {
+                todo!();
+            }
+
+            let rbbytes = rb.str_from_slice(buf);
 
-        let number_bytes_written = self
-            .inner
-            .funcall::<_, _, usize>("write", (rbbytes,))
-            .map_err(rberr_to_io_err)?;
+            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> {
-        self.inner
-            .funcall::<_, _, Value>("flush", ())
-            .map_err(rberr_to_io_err)?;
-
-        Ok(())
+        Self::flush(self)
     }
 }
 
 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 new_position = self
-            .inner
-            .funcall("seek", (offset, whence))
-            .map_err(rberr_to_io_err)?;
+            let inner = rb.get_inner(self.inner);
 
-        Ok(new_position)
+            inner
+                .funcall::<_, _, Value>("seek", (offset, whence))
+                .map_err(rberr_to_io_err)?;
+
+            inner.funcall("tell", ()).map_err(rberr_to_io_err)
+        })
     }
 }
 
@@ -131,7 +199,7 @@ impl FileLike for RbFileLikeObject {}
 
 pub enum EitherRustRubyFile {
     Rb(RbFileLikeObject),
-    Rust(File),
+    Rust(std::fs::File),
 }
 
 impl EitherRustRubyFile {
@@ -141,43 +209,85 @@ impl EitherRustRubyFile {
             EitherRustRubyFile::Rust(f) => Box::new(f),
         }
     }
+
+    fn into_scan_source_input(self) -> RubyScanSourceInput {
+        match self {
+            EitherRustRubyFile::Rb(f) => RubyScanSourceInput::Buffer(f.to_buffer()),
+            EitherRustRubyFile::Rust(f) => RubyScanSourceInput::File(f),
+        }
+    }
+
+    pub(crate) fn into_writeable(self) -> Writeable {
+        match self {
+            Self::Rb(f) => Writeable::Dyn(Box::new(f)),
+            Self::Rust(f) => Writeable::Local(f),
+        }
+    }
 }
 
 pub enum RubyScanSourceInput {
-    Buffer(bytes::Bytes),
-    Path(PathBuf),
-    #[allow(dead_code)]
-    File(File),
+    Buffer(Buffer<u8>),
+    Path(PlRefPath),
+    File(std::fs::File),
 }
 
-pub fn get_ruby_scan_source_input(rb_f: Value, write: bool) -> RbResult<RubyScanSourceInput> {
-    if let Ok(file_path) = PathBuf::try_convert(rb_f) {
-        // TODO resolve_homedir
-        Ok(RubyScanSourceInput::Path(file_path))
-    } else {
-        let f = RbFileLikeObject::with_requirements(rb_f, !write, write, !write)?;
-        Ok(RubyScanSourceInput::Buffer(f.as_bytes()))
+pub(crate) fn try_get_rbfile(
+    rb: &Ruby,
+    rb_f: Value,
+    write: bool,
+) -> RbResult<(EitherRustRubyFile, Option<PathBuf>)> {
+    RbFileLikeObject::ensure_requirements(rb_f, !write, write, !write)?;
+    let expects_str = false;
+    let has_flush = rb_f.respond_to("flush", false)?;
+    let f = RbFileLikeObject::new(rb_f, expects_str, has_flush);
+
+    // TODO move
+    if write {
+        start_background_ruby_thread(rb);
     }
+
+    Ok((EitherRustRubyFile::Rb(f), None))
+}
+
+pub fn get_ruby_scan_source_input(rb_f: Value, write: bool) -> RbResult<RubyScanSourceInput> {
+    Ruby::attach(|rb| {
+        if let Ok(s) = PathBuf::try_convert(rb_f) {
+            let file_path =
+                PlRefPath::try_from_path(resolve_homedir(&s).as_ref()).map_err(to_rb_err)?;
+
+            Ok(RubyScanSourceInput::Path(file_path))
+        } else {
+            Ok(try_get_rbfile(rb, rb_f, write)?.0.into_scan_source_input())
+        }
+    })
+}
+
+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), Some(file_path.into_owned())))
+        } 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))
-    } 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>> {
@@ -204,11 +314,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(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_buffer())), path)),
+        },
     }
 }

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) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
-    let e = dsl::sum_horizontal(exprs).map_err(RbPolarsErr::from)?;
+pub fn sum_horizontal(exprs: RArray, ignore_nulls: bool) -> RbResult<RbExpr> {
+    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) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
-    let e = dsl::mean_horizontal(exprs).map_err(RbPolarsErr::from)?;
+pub fn mean_horizontal(exprs: RArray, ignore_nulls: bool) -> RbResult<RbExpr> {
+    let exprs = exprs.to_exprs()?;
+    let e = dsl::mean_horizontal(exprs, ignore_nulls).map_err(RbPolarsErr::from)?;
     Ok(e.into())
 }

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

@@ -2,14 +2,13 @@ use polars::lazy::dsl;
 
 use crate::RbExpr;
 
-// TODO add to Ruby
 pub fn business_day_count(
     start: &RbExpr,
     end: &RbExpr,
     week_mask: [bool; 7],
-    holidays: Vec<i32>,
+    holidays: &RbExpr,
 ) -> RbExpr {
     let start = start.inner.clone();
     let end = end.inner.clone();
-    dsl::business_day_count(start, end, week_mask, holidays).into()
+    dsl::business_day_count(start, end, week_mask, holidays.inner.clone()).into()
 }

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

@@ -59,11 +59,12 @@ pub fn concat_df_diagonal(seq: RArray) -> RbResult<RbDataFrame> {
     Ok(df.into())
 }
 
-pub fn concat_df_horizontal(seq: RArray) -> RbResult<RbDataFrame> {
+pub fn concat_df_horizontal(seq: RArray, strict: bool) -> RbResult<RbDataFrame> {
     let mut dfs = Vec::new();
     for item in seq.into_iter() {
         dfs.push(get_df(item)?);
     }
-    let df = functions::concat_df_horizontal(&dfs, true).map_err(RbPolarsErr::from)?;
+    let df =
+        functions::concat_df_horizontal(&dfs, true, strict, false).map_err(RbPolarsErr::from)?;
     Ok(df.into())
 }

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

@@ -1,19 +1,80 @@
-use magnus::{RHash, Value};
+use std::io::BufReader;
+
+use magnus::{RHash, Ruby, Value};
+use polars::prelude::ArrowSchema;
+use polars::prelude::CloudScheme;
 
 use crate::conversion::Wrap;
-use crate::file::get_file_like;
-use crate::prelude::DataType;
+use crate::file::{EitherRustRubyFile, get_either_file};
+use crate::io::cloud_options::OptRbCloudOptions;
+use crate::ruby::gvl::GvlExt;
+use crate::ruby::utils::TryIntoValue;
 use crate::{RbPolarsErr, RbResult};
 
-pub fn read_ipc_schema(rb_f: Value) -> RbResult<RHash> {
-    use polars_core::export::arrow::io::ipc::read::read_file_metadata;
-    let mut r = get_file_like(rb_f, false)?;
-    let metadata = read_file_metadata(&mut r).map_err(RbPolarsErr::from)?;
+pub fn read_ipc_schema(rb: &Ruby, rb_f: Value) -> RbResult<RHash> {
+    use arrow::io::ipc::read::read_file_metadata;
+    let metadata = match get_either_file(rb_f, false)? {
+        EitherRustRubyFile::Rust(r) => {
+            read_file_metadata(&mut BufReader::new(r)).map_err(RbPolarsErr::from)?
+        }
+        EitherRustRubyFile::Rb(mut r) => read_file_metadata(&mut r).map_err(RbPolarsErr::from)?,
+    };
+
+    let dict = rb.hash_new();
+    fields_to_rbdict(&metadata.schema, &dict)?;
+    Ok(dict)
+}
+
+pub fn read_parquet_metadata(
+    rb: &Ruby,
+    rb_f: Value,
+    storage_options: OptRbCloudOptions,
+    credential_provider: Option<Value>,
+) -> RbResult<RHash> {
+    use std::io::Cursor;
+
+    use polars_parquet::read::read_metadata;
+    use polars_parquet::read::schema::read_custom_key_value_metadata;
+
+    use crate::file::{RubyScanSourceInput, get_ruby_scan_source_input};
 
-    let dict = RHash::new();
-    for field in metadata.schema.iter_values() {
-        let dt: Wrap<DataType> = Wrap((&field.dtype).into());
-        dict.aset(field.name.as_str(), dt)?;
+    let metadata = match get_ruby_scan_source_input(rb_f, false)? {
+        RubyScanSourceInput::Buffer(buf) => {
+            read_metadata(&mut Cursor::new(buf)).map_err(RbPolarsErr::from)?
+        }
+        RubyScanSourceInput::Path(p) => {
+            let cloud_options = storage_options.extract_opt_cloud_options(
+                CloudScheme::from_path(p.as_str()),
+                credential_provider,
+            )?;
+
+            if p.has_scheme() {
+                use polars::prelude::ParquetObjectStore;
+                use polars_error::PolarsResult;
+
+                rb.detach(|| {
+                    polars_core::runtime::ASYNC.block_on(async {
+                        let mut reader =
+                            ParquetObjectStore::from_uri(p, cloud_options.as_ref(), None).await?;
+                        let result = reader.get_metadata().await?;
+                        PolarsResult::Ok((**result).clone())
+                    })
+                })
+                .map_err(RbPolarsErr::from)?
+            } else {
+                let file = polars_utils::open_file(p.as_std_path()).map_err(RbPolarsErr::from)?;
+                read_metadata(&mut BufReader::new(file)).map_err(RbPolarsErr::from)?
+            }
+        }
+        RubyScanSourceInput::File(f) => {
+            read_metadata(&mut BufReader::new(f)).map_err(RbPolarsErr::from)?
+        }
+    };
+
+    let key_value_metadata = read_custom_key_value_metadata(metadata.key_value_metadata());
+    let dict = rb.hash_new();
+    for (key, value) in key_value_metadata.into_iter() {
+        dict.aset(key.as_str(), value.as_str())?;
     }
     Ok(dict)
 }
@@ -21,14 +82,25 @@ pub fn read_ipc_schema(rb_f: Value) -> RbResult<RHash> {
 pub fn read_parquet_schema(rb_f: Value) -> RbResult<RHash> {
     use polars_parquet::read::{infer_schema, read_metadata};
 
-    let mut r = get_file_like(rb_f, false)?;
-    let metadata = read_metadata(&mut r).map_err(RbPolarsErr::from)?;
+    let metadata = match get_either_file(rb_f, false)? {
+        EitherRustRubyFile::Rust(r) => {
+            read_metadata(&mut BufReader::new(r)).map_err(RbPolarsErr::from)?
+        }
+        EitherRustRubyFile::Rb(mut r) => read_metadata(&mut r).map_err(RbPolarsErr::from)?,
+    };
     let arrow_schema = infer_schema(&metadata).map_err(RbPolarsErr::from)?;
 
-    let dict = RHash::new();
-    for field in arrow_schema.iter_values() {
-        let dt: Wrap<DataType> = Wrap((&field.dtype).into());
-        dict.aset(field.name.as_str(), dt)?;
-    }
+    let ruby = Ruby::get_with(rb_f);
+    let dict = ruby.hash_new();
+    fields_to_rbdict(&arrow_schema, &dict)?;
     Ok(dict)
 }
+
+fn fields_to_rbdict(schema: &ArrowSchema, dict: &RHash) -> RbResult<()> {
+    let ruby = &Ruby::get_with(*dict);
+    for field in schema.iter_values() {
+        let dt = Wrap(polars::prelude::DataType::from_arrow_field(field));
+        dict.aset(field.name.as_str(), dt.try_into_value_with(ruby)?)?;
+    }
+    Ok(())
+}