polars-df 0.14.0 → 0.15.0

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

@@ -9,7 +9,6 @@ use polars_core::series::IsSorted;
 use crate::conversion::{parse_fill_null_strategy, Wrap};
 use crate::map::lazy::map_single;
 use crate::rb_exprs_to_exprs;
-use crate::utils::reinterpret;
 use crate::{RbExpr, RbResult};
 
 impl RbExpr {
@@ -165,7 +164,7 @@ impl RbExpr {
         self.inner.clone().implode().into()
     }
 
-    pub fn quantile(&self, quantile: &Self, interpolation: Wrap<QuantileInterpolOptions>) -> Self {
+    pub fn quantile(&self, quantile: &Self, interpolation: Wrap<QuantileMethod>) -> Self {
         self.inner
             .clone()
             .quantile(quantile.inner.clone(), interpolation.0)
@@ -463,14 +462,7 @@ impl RbExpr {
     }
 
     pub fn gather_every(&self, n: usize, offset: usize) -> Self {
-        self.clone()
-            .inner
-            .map(
-                move |s: Series| Ok(Some(s.gather_every(n, offset))),
-                GetOutput::same_type(),
-            )
-            .with_fmt("gather_every")
-            .into()
+        self.inner.clone().gather_every(n, offset).into()
     }
 
     pub fn tail(&self, n: Option<usize>) -> Self {
@@ -644,8 +636,16 @@ impl RbExpr {
         output_type: Option<Wrap<DataType>>,
         agg_list: bool,
         is_elementwise: bool,
+        returns_scalar: bool,
     ) -> Self {
-        map_single(self, lambda, output_type, agg_list, is_elementwise)
+        map_single(
+            self,
+            lambda,
+            output_type,
+            agg_list,
+            is_elementwise,
+            returns_scalar,
+        )
     }
 
     pub fn dot(&self, other: &Self) -> Self {
@@ -653,16 +653,7 @@ impl RbExpr {
     }
 
     pub fn reinterpret(&self, signed: bool) -> Self {
-        let function = move |s: Series| reinterpret(&s, signed).map(Some);
-        let dt = if signed {
-            DataType::Int64
-        } else {
-            DataType::UInt64
-        };
-        self.clone()
-            .inner
-            .map(function, GetOutput::from_type(dt))
-            .into()
+        self.inner.clone().reinterpret(signed).into()
     }
 
     pub fn mode(&self) -> Self {
@@ -717,7 +708,7 @@ impl RbExpr {
     }
 
     pub fn reshape(&self, dims: Vec<i64>) -> Self {
-        self.inner.clone().reshape(&dims, NestedType::Array).into()
+        self.inner.clone().reshape(&dims).into()
     }
 
     pub fn cum_count(&self, reverse: bool) -> Self {

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

@@ -1,4 +1,4 @@
-use magnus::Value;
+use magnus::{prelude::*, value::Opaque, Ruby, Value};
 use polars::lazy::dsl::lit;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
@@ -179,23 +179,30 @@ impl RbExpr {
     pub fn list_to_struct(
         &self,
         width_strat: Wrap<ListToStructWidthStrategy>,
-        _name_gen: Option<Value>,
+        name_gen: Option<Value>,
         upper_bound: usize,
     ) -> RbResult<Self> {
-        // TODO fix
-        let name_gen = None;
-        // let name_gen = name_gen.map(|lambda| {
-        //     Arc::new(move |idx: usize| {
-        //         let out: Value = lambda.funcall("call", (idx,)).unwrap();
-        //         String::try_convert(out).unwrap()
-        //     }) as NameGenerator
-        // });
+        let name_gen = name_gen.map(|lambda| {
+            let lambda = Opaque::from(lambda);
+            Arc::new(move |idx: usize| {
+                let lambda = Ruby::get().unwrap().get_inner(lambda);
+                let out: String = lambda.funcall("call", (idx,)).unwrap();
+                PlSmallStr::from_string(out)
+            });
+
+            // non-Ruby thread
+            todo!();
+        });
 
         Ok(self
             .inner
             .clone()
             .list()
-            .to_struct(width_strat.0, name_gen, upper_bound)
+            .to_struct(ListToStructArgs::InferWidth {
+                infer_field_strategy: width_strat.0,
+                get_index_name: name_gen,
+                max_fields: upper_bound,
+            })
             .into())
     }
 

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

@@ -1,5 +1,4 @@
 use polars::prelude::*;
-use std::any::Any;
 
 use crate::conversion::Wrap;
 use crate::RbExpr;
@@ -169,7 +168,7 @@ impl RbExpr {
             weights,
             min_periods,
             center,
-            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+            fn_params: Some(RollingFnParams::Var(RollingVarParams { ddof })),
         };
 
         self.inner.clone().rolling_std(options).into()
@@ -187,7 +186,7 @@ impl RbExpr {
             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>),
+            fn_params: Some(RollingFnParams::Var(RollingVarParams { ddof })),
         };
 
         self.inner
@@ -210,7 +209,7 @@ impl RbExpr {
             weights,
             min_periods,
             center,
-            fn_params: Some(Arc::new(RollingVarParams { ddof }) as Arc<dyn Any + Send + Sync>),
+            fn_params: Some(RollingFnParams::Var(RollingVarParams { ddof })),
         };
 
         self.inner.clone().rolling_var(options).into()
@@ -228,7 +227,7 @@ impl RbExpr {
             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>),
+            fn_params: Some(RollingFnParams::Var(RollingVarParams { ddof })),
         };
 
         self.inner
@@ -277,7 +276,7 @@ impl RbExpr {
     pub fn rolling_quantile(
         &self,
         quantile: f64,
-        interpolation: Wrap<QuantileInterpolOptions>,
+        interpolation: Wrap<QuantileMethod>,
         window_size: usize,
         weights: Option<Vec<f64>>,
         min_periods: Option<usize>,
@@ -302,7 +301,7 @@ impl RbExpr {
         &self,
         by: &RbExpr,
         quantile: f64,
-        interpolation: Wrap<QuantileInterpolOptions>,
+        interpolation: Wrap<QuantileMethod>,
         window_size: String,
         min_periods: usize,
         closed: Wrap<ClosedWindow>,

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

@@ -130,6 +130,11 @@ impl RbExpr {
         self.inner.clone().str().to_lowercase().into()
     }
 
+    // requires nightly
+    // pub fn str_to_titlecase(&self) -> Self {
+    //     self.inner.clone().str().to_titlecase().into()
+    // }
+
     pub fn str_len_bytes(&self) -> Self {
         self.inner.clone().str().len_bytes().into()
     }
@@ -200,51 +205,19 @@ impl RbExpr {
     }
 
     pub fn str_hex_encode(&self) -> Self {
-        self.clone()
-            .inner
-            .map(
-                move |s| s.str().map(|s| Some(s.hex_encode().into_series())),
-                GetOutput::same_type(),
-            )
-            .with_fmt("str.hex_encode")
-            .into()
+        self.inner.clone().str().hex_encode().into()
     }
 
     pub fn str_hex_decode(&self, strict: bool) -> Self {
-        self.clone()
-            .inner
-            .map(
-                move |s| s.str()?.hex_decode(strict).map(|s| Some(s.into_series())),
-                GetOutput::same_type(),
-            )
-            .with_fmt("str.hex_decode")
-            .into()
+        self.inner.clone().str().hex_decode(strict).into()
     }
 
     pub fn str_base64_encode(&self) -> Self {
-        self.clone()
-            .inner
-            .map(
-                move |s| s.str().map(|s| Some(s.base64_encode().into_series())),
-                GetOutput::same_type(),
-            )
-            .with_fmt("str.base64_encode")
-            .into()
+        self.inner.clone().str().base64_encode().into()
     }
 
     pub fn str_base64_decode(&self, strict: bool) -> Self {
-        self.clone()
-            .inner
-            .map(
-                move |s| {
-                    s.str()?
-                        .base64_decode(strict)
-                        .map(|s| Some(s.into_series()))
-                },
-                GetOutput::same_type(),
-            )
-            .with_fmt("str.base64_decode")
-            .into()
+        self.inner.clone().str().base64_decode(strict).into()
     }
 
     pub fn str_to_integer(&self, base: &Self, strict: bool) -> Self {

data/ext/polars/src/file.rs

@@ -1,6 +1,6 @@
 use std::fs::File;
 use std::io;
-use std::io::{BufReader, Cursor, Read, Seek, SeekFrom, Write};
+use std::io::{Cursor, Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 
 use magnus::{exception, prelude::*, Error, RString, Value};
@@ -24,9 +24,8 @@ impl RbFileLikeObject {
         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_bytes(&self) -> bytes::Bytes {
+        self.as_file_buffer().into_inner().into()
     }
 
     pub fn as_file_buffer(&self) -> Cursor<Vec<u8>> {
@@ -132,7 +131,33 @@ impl FileLike for RbFileLikeObject {}
 
 pub enum EitherRustRubyFile {
     Rb(RbFileLikeObject),
-    Rust(BufReader<File>),
+    Rust(File),
+}
+
+impl EitherRustRubyFile {
+    pub fn into_dyn(self) -> Box<dyn FileLike> {
+        match self {
+            EitherRustRubyFile::Rb(f) => Box::new(f),
+            EitherRustRubyFile::Rust(f) => Box::new(f),
+        }
+    }
+}
+
+pub enum RubyScanSourceInput {
+    Buffer(bytes::Bytes),
+    Path(PathBuf),
+    #[allow(dead_code)]
+    File(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()))
+    }
 }
 
 ///
@@ -144,12 +169,11 @@ pub fn get_either_file(rb_f: Value, truncate: bool) -> RbResult<EitherRustRubyFi
         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)?
+            File::create(file_path).map_err(RbPolarsErr::from)?
         } else {
             polars_utils::open_file(&file_path).map_err(RbPolarsErr::from)?
         };
-        let reader = BufReader::new(f);
-        Ok(EitherRustRubyFile::Rust(reader))
+        Ok(EitherRustRubyFile::Rust(f))
     } else {
         let f = RbFileLikeObject::with_requirements(rb_f, !truncate, truncate, !truncate)?;
         Ok(EitherRustRubyFile::Rb(f))
@@ -157,21 +181,34 @@ pub fn get_either_file(rb_f: Value, truncate: bool) -> RbResult<EitherRustRubyFi
 }
 
 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())),
-    }
+    Ok(get_either_file(f, truncate)?.into_dyn())
 }
 
-pub fn get_mmap_bytes_reader(rb_f: Value) -> RbResult<Box<dyn MmapBytesReader>> {
-    if let Ok(bytes) = rb_f.funcall::<_, _, RString>("read", ()) {
-        let bytes = unsafe { bytes.as_slice() };
-        // TODO avoid copy
-        Ok(Box::new(Cursor::new(bytes.to_vec())))
-    } else {
-        let p = PathBuf::try_convert(rb_f)?;
-        let f = File::open(p).map_err(|e| Error::new(exception::runtime_error(), e.to_string()))?;
-        Ok(Box::new(f))
+pub enum RbReadBytes {
+    Bytes(RString),
+    Other(Value),
+}
+
+pub fn read_if_bytesio(rb_f: Value) -> RbReadBytes {
+    rb_f.funcall("read", ())
+        .map(RbReadBytes::Bytes)
+        .unwrap_or(RbReadBytes::Other(rb_f))
+}
+
+pub fn get_mmap_bytes_reader<'a>(rb_f: &'a RbReadBytes) -> RbResult<Box<dyn MmapBytesReader + 'a>> {
+    get_mmap_bytes_reader_and_path(rb_f).map(|t| t.0)
+}
+
+pub fn get_mmap_bytes_reader_and_path<'a>(
+    rb_f: &'a RbReadBytes,
+) -> 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)))
+        }
     }
 }

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

@@ -0,0 +1,15 @@
+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>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    dsl::business_day_count(start, end, week_mask, holidays).into()
+}

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

@@ -94,10 +94,7 @@ pub fn col(name: String) -> RbExpr {
 }
 
 pub fn collect_all(lfs: RArray) -> RbResult<RArray> {
-    let lfs = lfs
-        .into_iter()
-        .map(<&RbLazyFrame>::try_convert)
-        .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
+    let lfs = lfs.typecheck::<Obj<RbLazyFrame>>()?;
 
     Ok(RArray::from_iter(lfs.iter().map(|lf| {
         let df = lf.ldf.borrow().clone().collect().unwrap();
@@ -173,8 +170,14 @@ pub fn cum_fold(
     let exprs = rb_exprs_to_exprs(exprs)?;
     let lambda = Opaque::from(lambda);
 
-    let func =
-        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
+    let func = move |a: Column, b: Column| {
+        binary_lambda(
+            Ruby::get().unwrap().get_inner(lambda),
+            a.take_materialized_series(),
+            b.take_materialized_series(),
+        )
+        .map(|v| v.map(Column::from))
+    };
     Ok(dsl::cum_fold_exprs(acc.inner.clone(), func, exprs, include_init).into())
 }
 
@@ -263,8 +266,14 @@ pub fn fold(acc: &RbExpr, lambda: Value, exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
     let lambda = Opaque::from(lambda);
 
-    let func =
-        move |a: Series, b: Series| binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b);
+    let func = move |a: Column, b: Column| {
+        binary_lambda(
+            Ruby::get().unwrap().get_inner(lambda),
+            a.take_materialized_series(),
+            b.take_materialized_series(),
+        )
+        .map(|v| v.map(Column::from))
+    };
     Ok(dsl::fold_exprs(acc.inner.clone(), func, exprs).into())
 }
 

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

@@ -1,4 +1,5 @@
 pub mod aggregation;
+pub mod business;
 pub mod eager;
 pub mod io;
 pub mod lazy;

data/ext/polars/src/interop/arrow/mod.rs

@@ -0,0 +1 @@
+pub mod to_ruby;

data/ext/polars/src/interop/arrow/to_ruby.rs

@@ -0,0 +1,83 @@
+use arrow::datatypes::ArrowDataType;
+use arrow::ffi;
+use magnus::{IntoValue, Value};
+use polars::datatypes::CompatLevel;
+use polars::frame::DataFrame;
+use polars::prelude::{ArrayRef, ArrowField, PlSmallStr, PolarsResult, SchemaExt};
+use polars::series::Series;
+use polars_core::utils::arrow;
+
+use crate::RbResult;
+
+#[magnus::wrap(class = "Polars::RbArrowArrayStream")]
+pub struct RbArrowArrayStream {
+    stream: ffi::ArrowArrayStream,
+}
+
+impl RbArrowArrayStream {
+    pub fn to_i(&self) -> usize {
+        (&self.stream as *const _) as usize
+    }
+}
+
+pub(crate) fn dataframe_to_stream(df: &DataFrame) -> RbResult<Value> {
+    let iter = Box::new(DataFrameStreamIterator::new(df));
+    let field = iter.field();
+    let stream = ffi::export_iterator(iter, field);
+    Ok(RbArrowArrayStream { stream }.into_value())
+}
+
+pub struct DataFrameStreamIterator {
+    columns: Vec<Series>,
+    dtype: ArrowDataType,
+    idx: usize,
+    n_chunks: usize,
+}
+
+impl DataFrameStreamIterator {
+    fn new(df: &DataFrame) -> Self {
+        let schema = df.schema().to_arrow(CompatLevel::newest());
+        let dtype = ArrowDataType::Struct(schema.into_iter_values().collect());
+
+        Self {
+            columns: df
+                .get_columns()
+                .iter()
+                .map(|v| v.as_materialized_series().clone())
+                .collect(),
+            dtype,
+            idx: 0,
+            n_chunks: df.n_chunks(),
+        }
+    }
+
+    fn field(&self) -> ArrowField {
+        ArrowField::new(PlSmallStr::EMPTY, self.dtype.clone(), false)
+    }
+}
+
+impl Iterator for DataFrameStreamIterator {
+    type Item = PolarsResult<ArrayRef>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.idx >= self.n_chunks {
+            None
+        } else {
+            // create a batch of the columns with the same chunk no.
+            let batch_cols = self
+                .columns
+                .iter()
+                .map(|s| s.to_arrow(self.idx, CompatLevel::newest()))
+                .collect::<Vec<_>>();
+            self.idx += 1;
+
+            let array = arrow::array::StructArray::new(
+                self.dtype.clone(),
+                batch_cols[0].len(),
+                batch_cols,
+                None,
+            );
+            Some(Ok(Box::new(array)))
+        }
+    }
+}

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

@@ -1 +1,2 @@
+pub mod arrow;
 pub mod numo;