polars-df 0.14.0 → 0.16.0

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

@@ -1,13 +1,15 @@
-use magnus::{prelude::*, RString, Value};
+use magnus::{prelude::*, Value};
 use polars::io::avro::AvroCompression;
 use polars::io::RowIndex;
 use polars::prelude::*;
-use std::io::{BufWriter, Cursor};
+use std::io::BufWriter;
 use std::num::NonZeroUsize;
 
 use super::*;
 use crate::conversion::*;
-use crate::file::{get_either_file, get_file_like, get_mmap_bytes_reader, EitherRustRubyFile};
+use crate::file::{
+    get_file_like, get_mmap_bytes_reader, get_mmap_bytes_reader_and_path, read_if_bytesio,
+};
 use crate::{RbPolarsErr, RbResult};
 
 impl RbDataFrame {
@@ -29,8 +31,7 @@ impl RbDataFrame {
         let n_threads = Option::<usize>::try_convert(arguments[12])?;
         let path = Option::<String>::try_convert(arguments[13])?;
         let overwrite_dtype = Option::<Vec<(String, Wrap<DataType>)>>::try_convert(arguments[14])?;
-        // TODO fix
-        let overwrite_dtype_slice = Option::<Vec<Wrap<DataType>>>::None; // Option::<Vec<Wrap<DataType>>>::try_convert(arguments[15])?;
+        let overwrite_dtype_slice = Option::<Vec<Wrap<DataType>>>::try_convert(arguments[15])?;
         let low_memory = bool::try_convert(arguments[16])?;
         let comment_prefix = Option::<String>::try_convert(arguments[17])?;
         let quote_char = Option::<String>::try_convert(arguments[18])?;
@@ -39,12 +40,11 @@ impl RbDataFrame {
         let try_parse_dates = bool::try_convert(arguments[21])?;
         let skip_rows_after_header = usize::try_convert(arguments[22])?;
         let row_index = Option::<(String, IdxSize)>::try_convert(arguments[23])?;
-        let sample_size = usize::try_convert(arguments[24])?;
-        let eol_char = String::try_convert(arguments[25])?;
-        let raise_if_empty = bool::try_convert(arguments[26])?;
-        let truncate_ragged_lines = bool::try_convert(arguments[27])?;
-        let decimal_comma = bool::try_convert(arguments[28])?;
-        let schema = Option::<Wrap<Schema>>::try_convert(arguments[29])?;
+        let eol_char = String::try_convert(arguments[24])?;
+        let raise_if_empty = bool::try_convert(arguments[25])?;
+        let truncate_ragged_lines = bool::try_convert(arguments[26])?;
+        let decimal_comma = bool::try_convert(arguments[27])?;
+        let schema = Option::<Wrap<Schema>>::try_convert(arguments[28])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
@@ -80,7 +80,8 @@ impl RbDataFrame {
                 .collect::<Vec<_>>()
         });
 
-        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+        let rb_f = read_if_bytesio(rb_f);
+        let mmap_bytes_r = get_mmap_bytes_reader(&rb_f)?;
         let df = CsvReadOptions::default()
             .with_path(path)
             .with_infer_schema_length(infer_schema_length)
@@ -99,7 +100,6 @@ impl RbDataFrame {
             .with_low_memory(low_memory)
             .with_skip_rows_after_header(skip_rows_after_header)
             .with_row_index(row_index)
-            .with_sample_size(sample_size)
             .with_raise_if_empty(raise_if_empty)
             .with_parse_options(
                 CsvParseOptions::default()
@@ -120,59 +120,14 @@ impl RbDataFrame {
         Ok(df.into())
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub fn read_parquet(
-        rb_f: Value,
-        columns: Option<Vec<String>>,
-        projection: Option<Vec<usize>>,
-        n_rows: Option<usize>,
-        parallel: Wrap<ParallelStrategy>,
-        row_index: Option<(String, IdxSize)>,
-        low_memory: bool,
-        use_statistics: bool,
-        rechunk: bool,
-    ) -> RbResult<Self> {
-        use EitherRustRubyFile::*;
-
-        let row_index = row_index.map(|(name, offset)| RowIndex {
-            name: name.into(),
-            offset,
-        });
-        let result = match get_either_file(rb_f, false)? {
-            Rb(f) => {
-                let buf = f.as_buffer();
-                ParquetReader::new(buf)
-                    .with_projection(projection)
-                    .with_columns(columns)
-                    .read_parallel(parallel.0)
-                    .with_slice(n_rows.map(|x| (0, x)))
-                    .with_row_index(row_index)
-                    .set_low_memory(low_memory)
-                    .use_statistics(use_statistics)
-                    .set_rechunk(rechunk)
-                    .finish()
-            }
-            Rust(f) => ParquetReader::new(f.into_inner())
-                .with_projection(projection)
-                .with_columns(columns)
-                .read_parallel(parallel.0)
-                .with_slice(n_rows.map(|x| (0, x)))
-                .with_row_index(row_index)
-                .use_statistics(use_statistics)
-                .set_rechunk(rechunk)
-                .finish(),
-        };
-        let df = result.map_err(RbPolarsErr::from)?;
-        Ok(RbDataFrame::new(df))
-    }
-
     pub fn read_json(
         rb_f: Value,
         infer_schema_length: Option<usize>,
         schema: Option<Wrap<Schema>>,
         schema_overrides: Option<Wrap<Schema>>,
     ) -> RbResult<Self> {
-        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+        let rb_f = read_if_bytesio(rb_f);
+        let mmap_bytes_r = get_mmap_bytes_reader(&rb_f)?;
 
         let mut builder = JsonReader::new(mmap_bytes_r)
             .with_json_format(JsonFormat::Json)
@@ -196,7 +151,8 @@ impl RbDataFrame {
         schema: Option<Wrap<Schema>>,
         schema_overrides: Option<Wrap<Schema>>,
     ) -> RbResult<Self> {
-        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+        let rb_f = read_if_bytesio(rb_f);
+        let mmap_bytes_r = get_mmap_bytes_reader(&rb_f)?;
 
         let mut builder = JsonReader::new(mmap_bytes_r)
             .with_json_format(JsonFormat::JsonLines)
@@ -212,7 +168,7 @@ impl RbDataFrame {
 
         let out = builder
             .finish()
-            .map_err(|e| RbPolarsErr::other(format!("{e}")))?;
+            .map_err(|e| RbPolarsErr::Other(format!("{e}")))?;
         Ok(out.into())
     }
 
@@ -222,16 +178,16 @@ impl RbDataFrame {
         projection: Option<Vec<usize>>,
         n_rows: Option<usize>,
         row_index: Option<(String, IdxSize)>,
-        _memory_map: bool,
+        memory_map: bool,
     ) -> RbResult<Self> {
         let row_index = row_index.map(|(name, offset)| RowIndex {
             name: name.into(),
             offset,
         });
-        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+        let rb_f = read_if_bytesio(rb_f);
+        let (mmap_bytes_r, mmap_path) = get_mmap_bytes_reader_and_path(&rb_f)?;
 
-        // TODO fix
-        let mmap_path = None;
+        let mmap_path = if memory_map { mmap_path } else { None };
         let df = IpcReader::new(mmap_bytes_r)
             .with_projection(projection)
             .with_columns(columns)
@@ -255,8 +211,8 @@ impl RbDataFrame {
             name: name.into(),
             offset,
         });
-        // rb_f = read_if_bytesio(rb_f);
-        let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
+        let rb_f = read_if_bytesio(rb_f);
+        let mmap_bytes_r = get_mmap_bytes_reader(&rb_f)?;
         let df = IpcStreamReader::new(mmap_bytes_r)
             .with_projection(projection)
             .with_columns(columns)
@@ -302,41 +258,19 @@ impl RbDataFrame {
     ) -> RbResult<()> {
         let batch_size = batch_size.0;
         let null = null_value.unwrap_or_default();
-
-        if let Ok(s) = String::try_convert(rb_f) {
-            let f = std::fs::File::create(s).unwrap();
-            // no need for a buffered writer, because the csv writer does internal buffering
-            CsvWriter::new(f)
-                .include_header(include_header)
-                .with_separator(separator)
-                .with_quote_char(quote_char)
-                .with_batch_size(batch_size)
-                .with_datetime_format(datetime_format)
-                .with_date_format(date_format)
-                .with_time_format(time_format)
-                .with_float_precision(float_precision)
-                .with_null_value(null)
-                .finish(&mut self.df.borrow_mut())
-                .map_err(RbPolarsErr::from)?;
-        } else {
-            let mut buf = Cursor::new(Vec::new());
-            CsvWriter::new(&mut buf)
-                .include_header(include_header)
-                .with_separator(separator)
-                .with_quote_char(quote_char)
-                .with_batch_size(batch_size)
-                .with_datetime_format(datetime_format)
-                .with_date_format(date_format)
-                .with_time_format(time_format)
-                .with_float_precision(float_precision)
-                .with_null_value(null)
-                .finish(&mut self.df.borrow_mut())
-                .map_err(RbPolarsErr::from)?;
-            // TODO less copying
-            let rb_str = RString::from_slice(&buf.into_inner());
-            rb_f.funcall::<_, _, Value>("write", (rb_str,))?;
-        }
-
+        let mut buf = get_file_like(rb_f, true)?;
+        CsvWriter::new(&mut buf)
+            .include_header(include_header)
+            .with_separator(separator)
+            .with_quote_char(quote_char)
+            .with_batch_size(batch_size)
+            .with_datetime_format(datetime_format)
+            .with_date_format(date_format)
+            .with_time_format(time_format)
+            .with_float_precision(float_precision)
+            .with_null_value(null)
+            .finish(&mut self.df.borrow_mut())
+            .map_err(RbPolarsErr::from)?;
         Ok(())
     }
 
@@ -351,26 +285,14 @@ impl RbDataFrame {
     ) -> RbResult<()> {
         let compression = parse_parquet_compression(&compression, compression_level)?;
 
-        if let Ok(s) = String::try_convert(rb_f) {
-            let f = std::fs::File::create(s).unwrap();
-            ParquetWriter::new(f)
-                .with_compression(compression)
-                .with_statistics(statistics.0)
-                .with_row_group_size(row_group_size)
-                .with_data_page_size(data_page_size)
-                .finish(&mut self.df.borrow_mut())
-                .map_err(RbPolarsErr::from)?;
-        } else {
-            let buf = get_file_like(rb_f, true)?;
-            ParquetWriter::new(buf)
-                .with_compression(compression)
-                .with_statistics(statistics.0)
-                .with_row_group_size(row_group_size)
-                .with_data_page_size(data_page_size)
-                .finish(&mut self.df.borrow_mut())
-                .map_err(RbPolarsErr::from)?;
-        }
-
+        let buf = get_file_like(rb_f, true)?;
+        ParquetWriter::new(buf)
+            .with_compression(compression)
+            .with_statistics(statistics.0)
+            .with_row_group_size(row_group_size)
+            .with_data_page_size(data_page_size)
+            .finish(&mut self.df.borrow_mut())
+            .map_err(RbPolarsErr::from)?;
         Ok(())
     }
 
@@ -386,7 +308,7 @@ impl RbDataFrame {
             (false, _) => serde_json::to_writer(file, &*self.df.borrow())
                 .map_err(|e| PolarsError::ComputeError(format!("{:?}", e).into())),
         };
-        r.map_err(|e| RbPolarsErr::other(format!("{:?}", e)))?;
+        r.map_err(|e| RbPolarsErr::Other(format!("{:?}", e)))?;
         Ok(())
     }
 
@@ -397,7 +319,7 @@ impl RbDataFrame {
             .with_json_format(JsonFormat::JsonLines)
             .finish(&mut self.df.borrow_mut());
 
-        r.map_err(|e| RbPolarsErr::other(format!("{:?}", e)))?;
+        r.map_err(|e| RbPolarsErr::Other(format!("{:?}", e)))?;
         Ok(())
     }
 
@@ -405,23 +327,24 @@ impl RbDataFrame {
         &self,
         rb_f: Value,
         compression: Wrap<Option<IpcCompression>>,
+        compat_level: RbCompatLevel,
+        cloud_options: Option<Vec<(String, String)>>,
+        retries: usize,
     ) -> RbResult<()> {
-        if let Ok(s) = String::try_convert(rb_f) {
-            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)?;
+        let cloud_options = if let Ok(path) = String::try_convert(rb_f) {
+            let cloud_options = parse_cloud_options(&path, cloud_options.unwrap_or_default())?;
+            Some(cloud_options.with_max_retries(retries))
         } else {
-            let mut buf = Cursor::new(Vec::new());
-            IpcWriter::new(&mut buf)
-                .with_compression(compression.0)
-                .finish(&mut self.df.borrow_mut())
-                .map_err(RbPolarsErr::from)?;
-            // TODO less copying
-            let rb_str = RString::from_slice(&buf.into_inner());
-            rb_f.funcall::<_, _, Value>("write", (rb_str,))?;
-        }
+            None
+        };
+
+        let f = crate::file::try_get_writeable(rb_f, cloud_options.as_ref())?;
+
+        IpcWriter::new(f)
+            .with_compression(compression.0)
+            .with_compat_level(compat_level.0)
+            .finish(&mut self.df.borrow_mut())
+            .map_err(RbPolarsErr::from)?;
         Ok(())
     }
 
@@ -429,21 +352,14 @@ impl RbDataFrame {
         &self,
         rb_f: Value,
         compression: Wrap<Option<IpcCompression>>,
+        compat_level: RbCompatLevel,
     ) -> RbResult<()> {
-        if let Ok(s) = String::try_convert(rb_f) {
-            let f = std::fs::File::create(s).unwrap();
-            IpcStreamWriter::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)?;
-
-            IpcStreamWriter::new(&mut buf)
-                .with_compression(compression.0)
-                .finish(&mut self.df.borrow_mut())
-                .map_err(RbPolarsErr::from)?;
-        }
+        let mut buf = get_file_like(rb_f, true)?;
+        IpcStreamWriter::new(&mut buf)
+            .with_compression(compression.0)
+            .with_compat_level(compat_level.0)
+            .finish(&mut self.df.borrow_mut())
+            .map_err(RbPolarsErr::from)?;
         Ok(())
     }
 
@@ -451,23 +367,15 @@ impl RbDataFrame {
         &self,
         rb_f: Value,
         compression: Wrap<Option<AvroCompression>>,
+        name: String,
     ) -> RbResult<()> {
         use polars::io::avro::AvroWriter;
-
-        if let Ok(s) = String::try_convert(rb_f) {
-            let f = std::fs::File::create(s).unwrap();
-            AvroWriter::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)?;
-            AvroWriter::new(&mut buf)
-                .with_compression(compression.0)
-                .finish(&mut self.df.borrow_mut())
-                .map_err(RbPolarsErr::from)?;
-        }
-
+        let mut buf = get_file_like(rb_f, true)?;
+        AvroWriter::new(&mut buf)
+            .with_compression(compression.0)
+            .with_name(name)
+            .finish(&mut self.df.borrow_mut())
+            .map_err(RbPolarsErr::from)?;
         Ok(())
     }
 }

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

@@ -2,6 +2,7 @@ mod construction;
 mod export;
 mod general;
 mod io;
+mod serde;
 
 use polars::prelude::*;
 use std::cell::RefCell;

data/ext/polars/src/dataframe/serde.rs

@@ -0,0 +1,15 @@
+use crate::exceptions::ComputeError;
+use crate::file::get_file_like;
+use crate::{RbDataFrame, RbResult};
+use magnus::Value;
+use std::io::BufWriter;
+
+impl RbDataFrame {
+    // TODO add to Ruby
+    pub fn serialize_json(&self, rb_f: Value) -> RbResult<()> {
+        let file = get_file_like(rb_f, true)?;
+        let writer = BufWriter::new(file);
+        serde_json::to_writer(writer, &self.df)
+            .map_err(|err| ComputeError::new_err(err.to_string()))
+    }
+}

data/ext/polars/src/error.rs

@@ -1,67 +1,50 @@
-use magnus::exception;
+use std::fmt::{Debug, Formatter};
+
 use magnus::Error;
 use polars::prelude::PolarsError;
 
+use crate::exceptions::{ComputeError, InvalidOperationError};
 use crate::rb_modules;
 
-pub struct RbPolarsErr {}
-
-impl RbPolarsErr {
-    // convert to Error instead of Self
-    pub fn from(e: PolarsError) -> Error {
-        match e {
-            PolarsError::ComputeError(err) => ComputeError::new_err(err.to_string()),
-            PolarsError::InvalidOperation(err) => InvalidOperationError::new_err(err.to_string()),
-            _ => Error::new(rb_modules::error(), e.to_string()),
-        }
-    }
-
-    pub fn io(e: std::io::Error) -> Error {
-        Error::new(rb_modules::error(), e.to_string())
-    }
-
-    pub fn other(message: String) -> Error {
-        Error::new(rb_modules::error(), message)
-    }
-}
-
-pub struct RbTypeError {}
-
-impl RbTypeError {
-    pub fn new_err(message: String) -> Error {
-        Error::new(exception::type_error(), message)
-    }
+pub enum RbPolarsErr {
+    Polars(PolarsError),
+    Other(String),
 }
 
-pub struct RbValueError {}
-
-impl RbValueError {
-    pub fn new_err(message: String) -> Error {
-        Error::new(exception::arg_error(), message)
+impl From<PolarsError> for RbPolarsErr {
+    fn from(err: PolarsError) -> Self {
+        RbPolarsErr::Polars(err)
     }
 }
 
-pub struct RbOverflowError {}
-
-impl RbOverflowError {
-    pub fn new_err(message: String) -> Error {
-        Error::new(exception::range_error(), message)
+impl From<std::io::Error> for RbPolarsErr {
+    fn from(value: std::io::Error) -> Self {
+        RbPolarsErr::Other(format!("{value:?}"))
     }
 }
 
-pub struct ComputeError {}
-
-impl ComputeError {
-    pub fn new_err(message: String) -> Error {
-        Error::new(rb_modules::compute_error(), message)
+impl From<RbPolarsErr> for Error {
+    fn from(err: RbPolarsErr) -> Self {
+        match err {
+            RbPolarsErr::Polars(err) => match err {
+                PolarsError::ComputeError(err) => ComputeError::new_err(err.to_string()),
+                PolarsError::InvalidOperation(err) => {
+                    InvalidOperationError::new_err(err.to_string())
+                }
+                _ => Error::new(rb_modules::error(), err.to_string()),
+            },
+            RbPolarsErr::Other(err) => Error::new(rb_modules::error(), err.to_string()),
+        }
     }
 }
 
-pub struct InvalidOperationError {}
-
-impl InvalidOperationError {
-    pub fn new_err(message: String) -> Error {
-        Error::new(rb_modules::invalid_operation_error(), message)
+impl Debug for RbPolarsErr {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        use RbPolarsErr::*;
+        match self {
+            Polars(err) => write!(f, "{err:?}"),
+            Other(err) => write!(f, "BindingsError: {err:?}"),
+        }
     }
 }
 

data/ext/polars/src/exceptions.rs

@@ -0,0 +1,24 @@
+use crate::rb_modules;
+use magnus::{exception, Error};
+use std::borrow::Cow;
+
+macro_rules! create_exception {
+    ($type:ident, $cls:expr) => {
+        pub struct $type {}
+
+        impl $type {
+            pub fn new_err<T>(message: T) -> Error
+            where
+                T: Into<Cow<'static, str>>,
+            {
+                Error::new($cls, message)
+            }
+        }
+    };
+}
+
+create_exception!(RbTypeError, exception::type_error());
+create_exception!(RbValueError, exception::arg_error());
+create_exception!(RbOverflowError, exception::range_error());
+create_exception!(ComputeError, rb_modules::compute_error());
+create_exception!(InvalidOperationError, rb_modules::invalid_operation_error());

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

@@ -1,5 +1,3 @@
-use polars::prelude::*;
-
 use crate::RbExpr;
 
 impl RbExpr {
@@ -28,54 +26,18 @@ impl RbExpr {
     }
 
     pub fn bin_hex_decode(&self, strict: bool) -> Self {
-        self.clone()
-            .inner
-            .map(
-                move |s| {
-                    s.binary()?
-                        .hex_decode(strict)
-                        .map(|s| Some(s.into_series()))
-                },
-                GetOutput::same_type(),
-            )
-            .with_fmt("bin.hex_decode")
-            .into()
+        self.inner.clone().binary().hex_decode(strict).into()
     }
 
     pub fn bin_base64_decode(&self, strict: bool) -> Self {
-        self.clone()
-            .inner
-            .map(
-                move |s| {
-                    s.binary()?
-                        .base64_decode(strict)
-                        .map(|s| Some(s.into_series()))
-                },
-                GetOutput::same_type(),
-            )
-            .with_fmt("bin.base64_decode")
-            .into()
+        self.inner.clone().binary().base64_decode(strict).into()
     }
 
     pub fn bin_hex_encode(&self) -> Self {
-        self.clone()
-            .inner
-            .map(
-                move |s| s.binary().map(|s| Some(s.hex_encode().into_series())),
-                GetOutput::same_type(),
-            )
-            .with_fmt("bin.hex_encode")
-            .into()
+        self.inner.clone().binary().hex_encode().into()
     }
 
     pub fn bin_base64_encode(&self) -> Self {
-        self.clone()
-            .inner
-            .map(
-                move |s| s.binary().map(|s| Some(s.base64_encode().into_series())),
-                GetOutput::same_type(),
-            )
-            .with_fmt("bin.base64_encode")
-            .into()
+        self.inner.clone().binary().base64_encode().into()
     }
 }

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

@@ -13,12 +13,13 @@ impl RbExpr {
     }
 
     pub fn dt_epoch_seconds(&self) -> Self {
-        self.clone()
-            .inner
+        self.inner
+            .clone()
             .map(
                 |s| {
-                    s.timestamp(TimeUnit::Milliseconds)
-                        .map(|ca| Some((ca / 1000).into_series()))
+                    s.take_materialized_series()
+                        .timestamp(TimeUnit::Milliseconds)
+                        .map(|ca| Some((ca / 1000).into_column()))
                 },
                 GetOutput::from_type(DataType::Int64),
             )