polars-df 0.17.1 → 0.18.0

data/ext/polars/Cargo.toml

@@ -1,10 +1,10 @@
 [package]
 name = "polars"
-version = "0.17.1"
+version = "0.18.0"
 license = "MIT"
 authors = ["Andrew Kane <andrew@ankane.org>"]
 edition = "2021"
-rust-version = "1.80.0"
+rust-version = "1.85.0"
 publish = false
 
 [lib]
@@ -12,22 +12,22 @@ crate-type = ["cdylib"]
 
 [dependencies]
 ahash = "0.8"
-arrow = { package = "polars-arrow", version = "=0.46.0" }
+arrow = { package = "polars-arrow", version = "=0.47.1" }
 bytes = "1"
 chrono = "0.4"
 either = "1.8"
 magnus = "0.7"
 num-traits = "0.2"
-polars-core = "=0.46.0"
-polars-plan = "=0.46.0"
-polars-parquet = "=0.46.0"
-polars-utils = "=0.46.0"
+polars-core = "=0.47.1"
+polars-plan = "=0.47.1"
+polars-parquet = "=0.47.1"
+polars-utils = "=0.47.1"
 rayon = "1.9"
 regex = "1"
 serde_json = "1"
 
 [dependencies.polars]
-version = "=0.46.0"
+version = "=0.47.1"
 features = [
     "abs",
     "approx_unique",

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

@@ -21,7 +21,7 @@ use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 use polars_core::utils::arrow::array::Array;
 use polars_core::utils::materialize_dyn_int;
-use polars_plan::plans::ScanSources;
+use polars_plan::dsl::ScanSources;
 use polars_utils::mmap::MemSlice;
 use polars_utils::total_ord::{TotalEq, TotalHash};
 
@@ -219,7 +219,7 @@ impl IntoValue for Wrap<DataType> {
                     .funcall::<_, _, Value>("new", (tu.to_ascii(),))
                     .unwrap()
             }
-            DataType::Object(_, _) => {
+            DataType::Object(_) => {
                 let class = pl.const_get::<_, Value>("Object").unwrap();
                 class.funcall("new", ()).unwrap()
             }
@@ -332,7 +332,7 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::Array" => DataType::Array(Box::new(DataType::Null), 0),
                 "Polars::Struct" => DataType::Struct(vec![]),
                 "Polars::Null" => DataType::Null,
-                "Polars::Object" => DataType::Object(OBJECT_NAME, None),
+                "Polars::Object" => DataType::Object(OBJECT_NAME),
                 "Polars::Unknown" => DataType::Unknown(Default::default()),
                 dt => {
                     return Err(RbValueError::new_err(format!(
@@ -408,7 +408,7 @@ impl TryConvert for Wrap<DataType> {
                     DataType::Struct(fields)
                 }
                 "Polars::Null" => DataType::Null,
-                "Object" => DataType::Object(OBJECT_NAME, None),
+                "Object" => DataType::Object(OBJECT_NAME),
                 "Polars::Unknown" => DataType::Unknown(Default::default()),
                 dt => {
                     return Err(RbTypeError::new_err(format!(
@@ -437,7 +437,7 @@ impl TryConvert for Wrap<DataType> {
                 "time" => DataType::Time,
                 "dur" => DataType::Duration(TimeUnit::Microseconds),
                 "f64" => DataType::Float64,
-                "obj" => DataType::Object(OBJECT_NAME, None),
+                "obj" => DataType::Object(OBJECT_NAME),
                 "list" => DataType::List(Box::new(DataType::Boolean)),
                 "null" => DataType::Null,
                 "unk" => DataType::Unknown(Default::default()),
@@ -761,6 +761,21 @@ impl TryConvert for Wrap<ClosedWindow> {
     }
 }
 
+impl TryConvert for Wrap<RoundMode> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match String::try_convert(ob)?.as_str() {
+            "half_to_even" => RoundMode::HalfToEven,
+            "half_away_from_zero" => RoundMode::HalfAwayFromZero,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "`mode` must be one of {{'half_to_even', 'half_away_from_zero'}}, got {v}",
+                )));
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<CsvEncoding> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match String::try_convert(ob)?.as_str() {

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

@@ -18,7 +18,7 @@ impl RbDataFrame {
                 .get_columns()
                 .iter()
                 .map(|s| match s.dtype() {
-                    DataType::Object(_, _) => {
+                    DataType::Object(_) => {
                         let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
                         obj.unwrap().to_value()
                     }
@@ -37,7 +37,7 @@ impl RbDataFrame {
                     .get_columns()
                     .iter()
                     .map(|s| match s.dtype() {
-                        DataType::Object(_, _) => {
+                        DataType::Object(_) => {
                             let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
                             obj.unwrap().to_value()
                         }

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

@@ -1,3 +1,5 @@
+use std::hash::BuildHasher;
+
 use either::Either;
 use magnus::{prelude::*, typed_data::Obj, IntoValue, RArray, Value};
 use polars::prelude::pivot::{pivot, pivot_stable};
@@ -494,7 +496,8 @@ impl RbDataFrame {
     }
 
     pub fn hash_rows(&self, k0: u64, k1: u64, k2: u64, k3: u64) -> RbResult<RbSeries> {
-        let hb = ahash::RandomState::with_seeds(k0, k1, k2, k3);
+        let seed = PlFixedStateQuality::default().hash_one((k0, k1, k2, k3));
+        let hb = PlSeedableRandomStateQuality::seed_from_u64(seed);
         let hash = self
             .df
             .borrow_mut()

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

@@ -371,14 +371,6 @@ impl RbExpr {
             .into())
     }
 
-    pub fn backward_fill(&self, limit: FillNullLimit) -> Self {
-        self.inner.clone().backward_fill(limit).into()
-    }
-
-    pub fn forward_fill(&self, limit: FillNullLimit) -> Self {
-        self.inner.clone().forward_fill(limit).into()
-    }
-
     pub fn shift(&self, n: &Self, fill_value: Option<&Self>) -> Self {
         let expr = self.inner.clone();
         let out = match fill_value {
@@ -497,8 +489,8 @@ impl RbExpr {
             .into()
     }
 
-    pub fn round(&self, decimals: u32) -> Self {
-        self.inner.clone().round(decimals).into()
+    pub fn round(&self, decimals: u32, mode: Wrap<RoundMode>) -> Self {
+        self.inner.clone().round(decimals, mode.0).into()
     }
 
     pub fn floor(&self) -> Self {
@@ -597,8 +589,11 @@ impl RbExpr {
         self.inner.clone().or(expr.inner.clone()).into()
     }
 
-    pub fn is_in(&self, expr: &Self) -> Self {
-        self.inner.clone().is_in(expr.inner.clone()).into()
+    pub fn is_in(&self, expr: &Self, nulls_equal: bool) -> Self {
+        self.inner
+            .clone()
+            .is_in(expr.inner.clone(), nulls_equal)
+            .into()
     }
 
     pub fn repeat_by(&self, by: &Self) -> Self {
@@ -698,8 +693,11 @@ impl RbExpr {
         self.inner.clone().rank(options, seed).into()
     }
 
-    pub fn diff(&self, n: i64, null_behavior: Wrap<NullBehavior>) -> Self {
-        self.inner.clone().diff(n, null_behavior.0).into()
+    pub fn diff(&self, n: &Self, null_behavior: Wrap<NullBehavior>) -> Self {
+        self.inner
+            .clone()
+            .diff(n.inner.clone(), null_behavior.0)
+            .into()
     }
 
     pub fn pct_change(&self, n: &Self) -> Self {

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

@@ -319,7 +319,22 @@ impl RbExpr {
             .into()
     }
 
-    pub fn rolling_skew(&self, window_size: usize, bias: bool) -> Self {
-        self.inner.clone().rolling_skew(window_size, bias).into()
+    pub fn rolling_skew(
+        &self,
+        window_size: usize,
+        bias: bool,
+        min_periods: Option<usize>,
+        center: bool,
+    ) -> Self {
+        let min_periods = min_periods.unwrap_or(window_size);
+        let options = RollingOptionsFixedWindow {
+            window_size,
+            weights: None,
+            min_periods,
+            center,
+            fn_params: Some(RollingFnParams::Skew { bias }),
+        };
+
+        self.inner.clone().rolling_skew(options).into()
     }
 }

data/ext/polars/src/file.rs

@@ -3,9 +3,12 @@ use std::io;
 use std::io::{Cursor, Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 
-use magnus::{exception, prelude::*, Error, RString, Value};
+use magnus::{exception, prelude::*, value::Opaque, Error, RString, Ruby, Value};
 use polars::io::cloud::CloudOptions;
 use polars::io::mmap::MmapBytesReader;
+use polars::prelude::file::DynWriteable;
+use polars::prelude::sync_on_close::SyncOnCloseType;
+use polars_utils::file::ClosableFile;
 use polars_utils::mmap::MemSlice;
 
 use crate::error::RbPolarsErr;
@@ -14,7 +17,22 @@ use crate::RbResult;
 
 #[derive(Clone)]
 pub struct RbFileLikeObject {
-    inner: Value,
+    inner: Opaque<Value>,
+}
+
+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(())
+    }
 }
 
 /// Wraps a `Value`, and implements read, seek, and write for it.
@@ -23,7 +41,9 @@ impl RbFileLikeObject {
     /// To assert the object has the required methods methods,
     /// instantiate it with `RbFileLikeObject::require`
     pub fn new(object: Value) -> Self {
-        RbFileLikeObject { inner: object }
+        RbFileLikeObject {
+            inner: object.into(),
+        }
     }
 
     pub fn as_bytes(&self) -> bytes::Bytes {
@@ -31,8 +51,9 @@ impl RbFileLikeObject {
     }
 
     pub fn as_file_buffer(&self) -> Cursor<Vec<u8>> {
-        let bytes = self
-            .inner
+        let bytes = Ruby::get()
+            .unwrap()
+            .get_inner(self.inner)
             .funcall::<_, _, RString>("read", ())
             .expect("no read method found");
 
@@ -77,8 +98,9 @@ fn rberr_to_io_err(e: Error) -> io::Error {
 
 impl Read for RbFileLikeObject {
     fn read(&mut self, mut buf: &mut [u8]) -> Result<usize, io::Error> {
-        let bytes = self
-            .inner
+        let bytes = Ruby::get()
+            .unwrap()
+            .get_inner(self.inner)
             .funcall::<_, _, RString>("read", (buf.len(),))
             .map_err(rberr_to_io_err)?;
 
@@ -92,8 +114,9 @@ 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
+        let number_bytes_written = Ruby::get()
+            .unwrap()
+            .get_inner(self.inner)
             .funcall::<_, _, usize>("write", (rbbytes,))
             .map_err(rberr_to_io_err)?;
 
@@ -101,7 +124,9 @@ impl Write for RbFileLikeObject {
     }
 
     fn flush(&mut self) -> Result<(), io::Error> {
-        self.inner
+        Ruby::get()
+            .unwrap()
+            .get_inner(self.inner)
             .funcall::<_, _, Value>("flush", ())
             .map_err(rberr_to_io_err)?;
 
@@ -117,8 +142,9 @@ impl Seek for RbFileLikeObject {
             SeekFrom::End(i) => (2, i),
         };
 
-        let new_position = self
-            .inner
+        let new_position = Ruby::get()
+            .unwrap()
+            .get_inner(self.inner)
             .funcall("seek", (offset, whence))
             .map_err(rberr_to_io_err)?;
 
@@ -129,11 +155,12 @@ impl Seek for RbFileLikeObject {
 pub trait FileLike: Read + Write + Seek {}
 
 impl FileLike for File {}
+impl FileLike for ClosableFile {}
 impl FileLike for RbFileLikeObject {}
 
 pub enum EitherRustRubyFile {
     Rb(RbFileLikeObject),
-    Rust(File),
+    Rust(ClosableFile),
 }
 
 impl EitherRustRubyFile {
@@ -144,6 +171,13 @@ impl EitherRustRubyFile {
         }
     }
 
+    pub(crate) fn into_writeable(self) -> Box<dyn DynWriteable> {
+        match self {
+            Self::Rb(f) => Box::new(f),
+            Self::Rust(f) => Box::new(f),
+        }
+    }
+
     pub fn into_dyn_writeable(self) -> Box<dyn Write> {
         match self {
             EitherRustRubyFile::Rb(f) => Box::new(f),
@@ -159,6 +193,14 @@ pub enum RubyScanSourceInput {
     File(File),
 }
 
+pub(crate) fn try_get_rbfile(
+    rb_f: Value,
+    write: bool,
+) -> RbResult<(EitherRustRubyFile, Option<PathBuf>)> {
+    let f = RbFileLikeObject::with_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(file_path) = PathBuf::try_convert(rb_f) {
         // TODO resolve_homedir
@@ -184,7 +226,7 @@ pub fn get_either_file(rb_f: Value, truncate: bool) -> RbResult<EitherRustRubyFi
         } else {
             polars_utils::open_file(&file_path).map_err(RbPolarsErr::from)?
         };
-        Ok(EitherRustRubyFile::Rust(f))
+        Ok(EitherRustRubyFile::Rust(f.into()))
     } else {
         let f = RbFileLikeObject::with_requirements(rb_f, !truncate, truncate, !truncate)?;
         Ok(EitherRustRubyFile::Rb(f))

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

@@ -278,7 +278,13 @@ pub fn first() -> RbExpr {
     dsl::first().into()
 }
 
-pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
+pub fn fold(
+    acc: &RbExpr,
+    lambda: Value,
+    exprs: RArray,
+    returns_scalar: bool,
+    return_dtype: Option<Wrap<DataType>>,
+) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
     let lambda = Opaque::from(lambda);
 
@@ -290,7 +296,14 @@ pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
         )
         .map(|v| v.map(Column::from))
     };
-    Ok(dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
+    Ok(dsl::fold_exprs(
+        acc.inner.clone(),
+        func,
+        exprs,
+        returns_scalar,
+        return_dtype.map(|w| w.0),
+    )
+    .into())
 }
 
 pub fn last() -> RbExpr {

data/ext/polars/src/lazyframe/general.rs

@@ -2,12 +2,13 @@ use magnus::{r_hash::ForEach, typed_data::Obj, IntoValue, RArray, RHash, TryConv
 use polars::io::{HiveOptions, RowIndex};
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
-use polars_plan::plans::ScanSources;
+use polars_plan::dsl::ScanSources;
 use std::cell::RefCell;
 use std::io::BufWriter;
 use std::num::NonZeroUsize;
 use std::path::PathBuf;
 
+use super::SinkTarget;
 use crate::conversion::*;
 use crate::expr::rb_exprs_to_exprs;
 use crate::file::get_file_like;
@@ -374,16 +375,16 @@ impl RbLazyFrame {
     #[allow(clippy::too_many_arguments)]
     pub fn sink_parquet(
         &self,
-        path: PathBuf,
+        target: SinkTarget,
         compression: String,
         compression_level: Option<i32>,
         statistics: Wrap<StatisticsOptions>,
         row_group_size: Option<usize>,
         data_page_size: Option<usize>,
-        maintain_order: bool,
         cloud_options: Option<Vec<(String, String)>>,
         retries: usize,
-    ) -> RbResult<()> {
+        sink_options: Wrap<SinkOptions>,
+    ) -> RbResult<RbLazyFrame> {
         let compression = parse_parquet_compression(&compression, compression_level)?;
 
         let options = ParquetWriteOptions {
@@ -391,48 +392,67 @@ impl RbLazyFrame {
             statistics: statistics.0,
             row_group_size,
             data_page_size,
-            maintain_order,
         };
 
-        let cloud_options = {
-            let cloud_options =
-                parse_cloud_options(path.to_str().unwrap(), cloud_options.unwrap_or_default())?;
-            Some(cloud_options.with_max_retries(retries))
+        let cloud_options = match target.base_path() {
+            None => None,
+            Some(base_path) => {
+                let cloud_options = parse_cloud_options(
+                    base_path.to_str().unwrap(),
+                    cloud_options.unwrap_or_default(),
+                )?;
+                Some(cloud_options.with_max_retries(retries))
+            }
         };
 
         let ldf = self.ldf.borrow().clone();
-        ldf.sink_parquet(&path, options, cloud_options)
-            .map_err(RbPolarsErr::from)?;
-        Ok(())
+        match target {
+            SinkTarget::File(target) => {
+                ldf.sink_parquet(target, options, cloud_options, sink_options.0)
+            }
+        }
+        .map_err(RbPolarsErr::from)
+        .map(Into::into)
+        .map_err(Into::into)
     }
 
     pub fn sink_ipc(
         &self,
-        path: PathBuf,
+        target: SinkTarget,
         compression: Option<Wrap<IpcCompression>>,
-        maintain_order: bool,
         cloud_options: Option<Vec<(String, String)>>,
         retries: usize,
-    ) -> RbResult<()> {
+        sink_options: Wrap<SinkOptions>,
+    ) -> RbResult<RbLazyFrame> {
         let options = IpcWriterOptions {
             compression: compression.map(|c| c.0),
-            maintain_order,
+            ..Default::default()
         };
 
-        let cloud_options = {
-            let cloud_options =
-                parse_cloud_options(path.to_str().unwrap(), cloud_options.unwrap_or_default())?;
-            Some(cloud_options.with_max_retries(retries))
+        let cloud_options = match target.base_path() {
+            None => None,
+            Some(base_path) => {
+                let cloud_options = parse_cloud_options(
+                    base_path.to_str().unwrap(),
+                    cloud_options.unwrap_or_default(),
+                )?;
+                Some(cloud_options.with_max_retries(retries))
+            }
         };
 
         let ldf = self.ldf.borrow().clone();
-        ldf.sink_ipc(&path, options, cloud_options)
-            .map_err(RbPolarsErr::from)?;
-        Ok(())
+        match target {
+            SinkTarget::File(target) => {
+                ldf.sink_ipc(target, options, cloud_options, sink_options.0)
+            }
+        }
+        .map_err(RbPolarsErr::from)
+        .map(Into::into)
+        .map_err(Into::into)
     }
 
-    pub fn sink_csv(&self, arguments: &[Value]) -> RbResult<()> {
-        let path = PathBuf::try_convert(arguments[0])?;
+    pub fn sink_csv(&self, arguments: &[Value]) -> RbResult<RbLazyFrame> {
+        let target = SinkTarget::try_convert(arguments[0])?;
         let include_bom = bool::try_convert(arguments[1])?;
         let include_header = bool::try_convert(arguments[2])?;
         let separator = u8::try_convert(arguments[3])?;
@@ -446,9 +466,9 @@ impl RbLazyFrame {
         let float_precision = Option::<usize>::try_convert(arguments[11])?;
         let null_value = Option::<String>::try_convert(arguments[12])?;
         let quote_style = Option::<Wrap<QuoteStyle>>::try_convert(arguments[13])?;
-        let maintain_order = bool::try_convert(arguments[14])?;
-        let cloud_options = Option::<Vec<(String, String)>>::try_convert(arguments[15])?;
-        let retries = usize::try_convert(arguments[16])?;
+        let cloud_options = Option::<Vec<(String, String)>>::try_convert(arguments[14])?;
+        let retries = usize::try_convert(arguments[15])?;
+        let sink_options = Wrap::<SinkOptions>::try_convert(arguments[16])?;
 
         let quote_style = quote_style.map_or(QuoteStyle::default(), |wrap| wrap.0);
         let null_value = null_value.unwrap_or(SerializeOptions::default().null);
@@ -469,42 +489,59 @@ impl RbLazyFrame {
         let options = CsvWriterOptions {
             include_bom,
             include_header,
-            maintain_order,
             batch_size: batch_size.0,
             serialize_options,
         };
 
-        let cloud_options = {
-            let cloud_options =
-                parse_cloud_options(path.to_str().unwrap(), cloud_options.unwrap_or_default())?;
-            Some(cloud_options.with_max_retries(retries))
+        let cloud_options = match target.base_path() {
+            None => None,
+            Some(base_path) => {
+                let cloud_options = parse_cloud_options(
+                    base_path.to_str().unwrap(),
+                    cloud_options.unwrap_or_default(),
+                )?;
+                Some(cloud_options.with_max_retries(retries))
+            }
         };
 
         let ldf = self.ldf.borrow().clone();
-        ldf.sink_csv(&path, options, cloud_options)
-            .map_err(RbPolarsErr::from)?;
-        Ok(())
+        match target {
+            SinkTarget::File(target) => {
+                ldf.sink_csv(target, options, cloud_options, sink_options.0)
+            }
+        }
+        .map_err(RbPolarsErr::from)
+        .map(Into::into)
+        .map_err(Into::into)
     }
 
     pub fn sink_json(
         &self,
-        path: PathBuf,
-        maintain_order: bool,
+        target: SinkTarget,
         cloud_options: Option<Vec<(String, String)>>,
         retries: usize,
-    ) -> RbResult<()> {
-        let options = JsonWriterOptions { maintain_order };
-
-        let cloud_options = {
-            let cloud_options =
-                parse_cloud_options(path.to_str().unwrap(), cloud_options.unwrap_or_default())?;
-            Some(cloud_options.with_max_retries(retries))
+        sink_options: Wrap<SinkOptions>,
+    ) -> RbResult<RbLazyFrame> {
+        let options = JsonWriterOptions {};
+
+        let cloud_options = match target.base_path() {
+            None => None,
+            Some(base_path) => {
+                let cloud_options = parse_cloud_options(
+                    base_path.to_str().unwrap(),
+                    cloud_options.unwrap_or_default(),
+                )?;
+                Some(cloud_options.with_max_retries(retries))
+            }
         };
 
         let ldf = self.ldf.borrow().clone();
-        ldf.sink_json(&path, options, cloud_options)
-            .map_err(RbPolarsErr::from)?;
-        Ok(())
+        match target {
+            SinkTarget::File(path) => ldf.sink_json(path, options, cloud_options, sink_options.0),
+        }
+        .map_err(RbPolarsErr::from)
+        .map(Into::into)
+        .map_err(Into::into)
     }
 
     pub fn fetch(&self, n_rows: usize) -> RbResult<RbDataFrame> {

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

@@ -1,7 +1,9 @@
 mod general;
 mod serde;
+mod sink;
 
 use polars::lazy::frame::LazyFrame;
+pub use sink::SinkTarget;
 use std::cell::RefCell;
 
 #[magnus::wrap(class = "Polars::RbLazyFrame")]