polars-df 0.14.0 → 0.16.0

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

@@ -2,25 +2,30 @@ pub(crate) mod any_value;
 mod chunked_array;
 
 use std::fmt::{Debug, Display, Formatter};
+use std::fs::File;
 use std::hash::{Hash, Hasher};
 use std::num::NonZeroUsize;
+use std::path::PathBuf;
 
 use magnus::{
-    class, exception, prelude::*, r_hash::ForEach, value::Opaque, IntoValue, Module, RArray, RHash,
-    Ruby, Symbol, TryConvert, Value,
+    class, exception, prelude::*, r_hash::ForEach, try_convert::TryConvertOwned, value::Opaque,
+    IntoValue, Module, RArray, RHash, Ruby, Symbol, TryConvert, Value,
 };
 use polars::chunked_array::object::PolarsObjectSafe;
 use polars::chunked_array::ops::{FillNullLimit, FillNullStrategy};
 use polars::datatypes::AnyValue;
 use polars::frame::row::Row;
-use polars::frame::NullStrategy;
 use polars::io::avro::AvroCompression;
+use polars::io::cloud::CloudOptions;
 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_utils::mmap::MemSlice;
 use polars_utils::total_ord::{TotalEq, TotalHash};
 
+use crate::file::{get_ruby_scan_source_input, RubyScanSourceInput};
 use crate::object::OBJECT_NAME;
 use crate::rb_modules::series;
 use crate::{RbDataFrame, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbTypeError, RbValueError};
@@ -104,9 +109,10 @@ impl TryConvert for Wrap<NullValues> {
                     .collect(),
             )))
         } else {
-            Err(RbPolarsErr::other(
-                "could not extract value from null_values argument".into(),
-            ))
+            Err(
+                RbPolarsErr::Other("could not extract value from null_values argument".into())
+                    .into(),
+            )
         }
     }
 }
@@ -298,29 +304,31 @@ impl TryConvert for Wrap<DataType> {
         let dtype = if ob.is_kind_of(class::class()) {
             let name = ob.funcall::<_, _, String>("name", ())?;
             match name.as_str() {
-                "Polars::UInt8" => DataType::UInt8,
-                "Polars::UInt16" => DataType::UInt16,
-                "Polars::UInt32" => DataType::UInt32,
-                "Polars::UInt64" => DataType::UInt64,
                 "Polars::Int8" => DataType::Int8,
                 "Polars::Int16" => DataType::Int16,
                 "Polars::Int32" => DataType::Int32,
                 "Polars::Int64" => DataType::Int64,
+                "Polars::UInt8" => DataType::UInt8,
+                "Polars::UInt16" => DataType::UInt16,
+                "Polars::UInt32" => DataType::UInt32,
+                "Polars::UInt64" => DataType::UInt64,
+                "Polars::Float32" => DataType::Float32,
+                "Polars::Float64" => DataType::Float64,
+                "Polars::Boolean" => DataType::Boolean,
                 "Polars::String" => DataType::String,
                 "Polars::Binary" => DataType::Binary,
-                "Polars::Boolean" => DataType::Boolean,
                 "Polars::Categorical" => DataType::Categorical(None, Default::default()),
                 "Polars::Enum" => DataType::Enum(None, Default::default()),
                 "Polars::Date" => DataType::Date,
-                "Polars::Datetime" => DataType::Datetime(TimeUnit::Microseconds, None),
                 "Polars::Time" => DataType::Time,
+                "Polars::Datetime" => DataType::Datetime(TimeUnit::Microseconds, None),
                 "Polars::Duration" => DataType::Duration(TimeUnit::Microseconds),
                 "Polars::Decimal" => DataType::Decimal(None, None),
-                "Polars::Float32" => DataType::Float32,
-                "Polars::Float64" => DataType::Float64,
-                "Polars::Object" => DataType::Object(OBJECT_NAME, None),
                 "Polars::List" => DataType::List(Box::new(DataType::Null)),
+                "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::Unknown" => DataType::Unknown(Default::default()),
                 dt => {
                     return Err(RbValueError::new_err(format!(
@@ -328,7 +336,6 @@ impl TryConvert for Wrap<DataType> {
                     )))
                 }
             }
-        // TODO improve
         } else if String::try_convert(ob).is_err() {
             let name = unsafe { ob.class().name() }.into_owned();
             match name.as_str() {
@@ -340,9 +347,11 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::UInt16" => DataType::UInt16,
                 "Polars::UInt32" => DataType::UInt32,
                 "Polars::UInt64" => DataType::UInt64,
+                "Polars::Float32" => DataType::Float32,
+                "Polars::Float64" => DataType::Float64,
+                "Polars::Boolean" => DataType::Boolean,
                 "Polars::String" => DataType::String,
                 "Polars::Binary" => DataType::Binary,
-                "Polars::Boolean" => DataType::Boolean,
                 "Polars::Categorical" => {
                     let ordering = ob
                         .funcall::<_, _, Wrap<CategoricalOrdering>>("ordering", ())?
@@ -358,21 +367,17 @@ impl TryConvert for Wrap<DataType> {
                 }
                 "Polars::Date" => DataType::Date,
                 "Polars::Time" => DataType::Time,
-                "Polars::Float32" => DataType::Float32,
-                "Polars::Float64" => DataType::Float64,
-                "Polars::Null" => DataType::Null,
-                "Polars::Unknown" => DataType::Unknown(Default::default()),
-                "Polars::Duration" => {
-                    let time_unit: Value = ob.funcall("time_unit", ()).unwrap();
-                    let time_unit = Wrap::<TimeUnit>::try_convert(time_unit)?.0;
-                    DataType::Duration(time_unit)
-                }
                 "Polars::Datetime" => {
                     let time_unit: Value = ob.funcall("time_unit", ()).unwrap();
                     let time_unit = Wrap::<TimeUnit>::try_convert(time_unit)?.0;
                     let time_zone: Option<String> = ob.funcall("time_zone", ())?;
                     DataType::Datetime(time_unit, time_zone.as_deref().map(|x| x.into()))
                 }
+                "Polars::Duration" => {
+                    let time_unit: Value = ob.funcall("time_unit", ()).unwrap();
+                    let time_unit = Wrap::<TimeUnit>::try_convert(time_unit)?.0;
+                    DataType::Duration(time_unit)
+                }
                 "Polars::Decimal" => {
                     let precision = ob.funcall("precision", ())?;
                     let scale = ob.funcall("scale", ())?;
@@ -383,6 +388,13 @@ impl TryConvert for Wrap<DataType> {
                     let inner = Wrap::<DataType>::try_convert(inner)?;
                     DataType::List(Box::new(inner.0))
                 }
+                "Polars::Array" => {
+                    let inner: Value = ob.funcall("inner", ()).unwrap();
+                    let size: Value = ob.funcall("size", ()).unwrap();
+                    let inner = Wrap::<DataType>::try_convert(inner)?;
+                    let size = usize::try_convert(size)?;
+                    DataType::Array(Box::new(inner.0), size)
+                }
                 "Polars::Struct" => {
                     let arr: RArray = ob.funcall("fields", ())?;
                     let mut fields = Vec::with_capacity(arr.len());
@@ -391,6 +403,9 @@ impl TryConvert for Wrap<DataType> {
                     }
                     DataType::Struct(fields)
                 }
+                "Polars::Null" => DataType::Null,
+                "Object" => DataType::Object(OBJECT_NAME, None),
+                "Polars::Unknown" => DataType::Unknown(Default::default()),
                 dt => {
                     return Err(RbTypeError::new_err(format!(
                         "A {dt} object is not a correct polars DataType. \
@@ -434,6 +449,8 @@ impl TryConvert for Wrap<DataType> {
     }
 }
 
+unsafe impl TryConvertOwned for Wrap<DataType> {}
+
 impl TryConvert for Wrap<StatisticsOptions> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let mut statistics = StatisticsOptions::empty();
@@ -452,8 +469,7 @@ impl TryConvert for Wrap<StatisticsOptions> {
                 }
             }
             Ok(ForEach::Continue)
-        })
-        .unwrap();
+        })?;
 
         Ok(Wrap(statistics))
     }
@@ -478,13 +494,75 @@ impl TryConvert for Wrap<Schema> {
         dict.foreach(|key: String, val: Wrap<DataType>| {
             schema.push(Ok(Field::new((&*key).into(), val.0)));
             Ok(ForEach::Continue)
-        })
-        .unwrap();
+        })?;
 
         Ok(Wrap(schema.into_iter().collect::<RbResult<Schema>>()?))
     }
 }
 
+impl TryConvert for Wrap<ScanSources> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let list = RArray::try_convert(ob)?;
+
+        if list.is_empty() {
+            return Ok(Wrap(ScanSources::default()));
+        }
+
+        enum MutableSources {
+            Paths(Vec<PathBuf>),
+            Files(Vec<File>),
+            Buffers(Vec<MemSlice>),
+        }
+
+        let num_items = list.len();
+        let mut iter = list
+            .into_iter()
+            .map(|val| get_ruby_scan_source_input(val, false));
+
+        let Some(first) = iter.next() else {
+            return Ok(Wrap(ScanSources::default()));
+        };
+
+        let mut sources = match first? {
+            RubyScanSourceInput::Path(path) => {
+                let mut sources = Vec::with_capacity(num_items);
+                sources.push(path);
+                MutableSources::Paths(sources)
+            }
+            RubyScanSourceInput::File(file) => {
+                let mut sources = Vec::with_capacity(num_items);
+                sources.push(file);
+                MutableSources::Files(sources)
+            }
+            RubyScanSourceInput::Buffer(buffer) => {
+                let mut sources = Vec::with_capacity(num_items);
+                sources.push(buffer);
+                MutableSources::Buffers(sources)
+            }
+        };
+
+        for source in iter {
+            match (&mut sources, source?) {
+                (MutableSources::Paths(v), RubyScanSourceInput::Path(p)) => v.push(p),
+                (MutableSources::Files(v), RubyScanSourceInput::File(f)) => v.push(f),
+                (MutableSources::Buffers(v), RubyScanSourceInput::Buffer(f)) => v.push(f),
+                _ => {
+                    return Err(RbTypeError::new_err(
+                        "Cannot combine in-memory bytes, paths and files for scan sources"
+                            .to_string(),
+                    ))
+                }
+            }
+        }
+
+        Ok(Wrap(match sources {
+            MutableSources::Paths(i) => ScanSources::Paths(i.into()),
+            MutableSources::Files(i) => ScanSources::Files(i.into()),
+            MutableSources::Buffers(i) => ScanSources::Buffers(i.into()),
+        }))
+    }
+}
+
 #[derive(Clone)]
 pub struct ObjectValue {
     pub inner: Opaque<Value>,
@@ -493,7 +571,7 @@ pub struct ObjectValue {
 impl Debug for ObjectValue {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("ObjectValue")
-            .field("inner", &self.to_object())
+            .field("inner", &self.to_value())
             .finish()
     }
 }
@@ -501,7 +579,7 @@ impl Debug for ObjectValue {
 impl Hash for ObjectValue {
     fn hash<H: Hasher>(&self, state: &mut H) {
         let h = self
-            .to_object()
+            .to_value()
             .funcall::<_, _, isize>("hash", ())
             .expect("should be hashable");
         state.write_isize(h)
@@ -512,7 +590,7 @@ impl Eq for ObjectValue {}
 
 impl PartialEq for ObjectValue {
     fn eq(&self, other: &Self) -> bool {
-        self.to_object().eql(other.to_object()).unwrap_or(false)
+        self.to_value().eql(other.to_value()).unwrap_or(false)
     }
 }
 
@@ -533,7 +611,7 @@ impl TotalHash for ObjectValue {
 
 impl Display for ObjectValue {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.to_object())
+        write!(f, "{}", self.to_value())
     }
 }
 
@@ -561,16 +639,15 @@ impl From<&dyn PolarsObjectSafe> for &ObjectValue {
     }
 }
 
-// TODO remove
 impl ObjectValue {
-    pub fn to_object(&self) -> Value {
-        Ruby::get().unwrap().get_inner(self.inner)
+    pub fn to_value(&self) -> Value {
+        self.clone().into_value()
     }
 }
 
 impl IntoValue for ObjectValue {
-    fn into_value_with(self, _: &Ruby) -> Value {
-        self.to_object()
+    fn into_value_with(self, ruby: &Ruby) -> Value {
+        ruby.get_inner(self.inner)
     }
 }
 
@@ -587,10 +664,10 @@ impl TryConvert for Wrap<AsofStrategy> {
         let parsed = match String::try_convert(ob)?.as_str() {
             "backward" => AsofStrategy::Backward,
             "forward" => AsofStrategy::Forward,
+            "nearest" => AsofStrategy::Nearest,
             v => {
                 return Err(RbValueError::new_err(format!(
-                    "strategy must be one of {{'backward', 'forward'}}, got {}",
-                    v
+                    "asof `strategy` must be one of {{'backward', 'forward', 'nearest'}}, got {v}",
                 )))
             }
         };
@@ -830,14 +907,14 @@ impl TryConvert for Wrap<ParallelStrategy> {
     }
 }
 
-impl TryConvert for Wrap<QuantileInterpolOptions> {
+impl TryConvert for Wrap<QuantileMethod> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match String::try_convert(ob)?.as_str() {
-            "lower" => QuantileInterpolOptions::Lower,
-            "higher" => QuantileInterpolOptions::Higher,
-            "nearest" => QuantileInterpolOptions::Nearest,
-            "linear" => QuantileInterpolOptions::Linear,
-            "midpoint" => QuantileInterpolOptions::Midpoint,
+            "lower" => QuantileMethod::Lower,
+            "higher" => QuantileMethod::Higher,
+            "nearest" => QuantileMethod::Nearest,
+            "linear" => QuantileMethod::Linear,
+            "midpoint" => QuantileMethod::Midpoint,
             v => {
                 return Err(RbValueError::new_err(format!(
                     "interpolation must be one of {{'lower', 'higher', 'nearest', 'linear', 'midpoint'}}, got {}",
@@ -1001,6 +1078,11 @@ impl TryConvert for Wrap<QuoteStyle> {
     }
 }
 
+pub(crate) fn parse_cloud_options(uri: &str, kv: Vec<(String, String)>) -> RbResult<CloudOptions> {
+    let out = CloudOptions::from_untyped_config(uri, kv).map_err(RbPolarsErr::from)?;
+    Ok(out)
+}
+
 pub fn parse_fill_null_strategy(
     strategy: &str,
     limit: FillNullLimit,
@@ -1071,7 +1153,7 @@ impl TryConvert for Wrap<NonZeroUsize> {
         let v = usize::try_convert(ob)?;
         NonZeroUsize::new(v)
             .map(Wrap)
-            .ok_or(RbValueError::new_err("must be non-zero".into()))
+            .ok_or(RbValueError::new_err("must be non-zero"))
     }
 }
 
@@ -1085,3 +1167,28 @@ where
         .map(|s| PlSmallStr::from_str(s.as_ref()))
         .collect()
 }
+
+#[derive(Debug, Copy, Clone)]
+pub struct RbCompatLevel(pub CompatLevel);
+
+impl TryConvert for RbCompatLevel {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        Ok(RbCompatLevel(if let Ok(level) = u16::try_convert(ob) {
+            if let Ok(compat_level) = CompatLevel::with_level(level) {
+                compat_level
+            } else {
+                return Err(RbValueError::new_err("invalid compat level".to_string()));
+            }
+        } else if let Ok(future) = bool::try_convert(ob) {
+            if future {
+                CompatLevel::newest()
+            } else {
+                CompatLevel::oldest()
+            }
+        } else {
+            return Err(RbTypeError::new_err(
+                "'compat_level' argument accepts int or bool".to_string(),
+            ));
+        }))
+    }
+}

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

@@ -54,9 +54,6 @@ fn finish_from_rows(
     schema_overrides: Option<Schema>,
     infer_schema_length: Option<usize>,
 ) -> RbResult<RbDataFrame> {
-    // Object builder must be registered
-    crate::on_startup::register_object_builder();
-
     let mut schema = if let Some(mut schema) = schema {
         resolve_schema_overrides(&mut schema, schema_overrides);
         update_schema_from_rows(&mut schema, &rows, infer_schema_length)?;

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

@@ -2,6 +2,8 @@ use magnus::{prelude::*, IntoValue, RArray, Value};
 
 use super::*;
 use crate::conversion::{ObjectValue, Wrap};
+use crate::interop::arrow::to_ruby::dataframe_to_stream;
+use crate::RbResult;
 
 impl RbDataFrame {
     pub fn row_tuple(&self, idx: i64) -> Value {
@@ -18,7 +20,7 @@ impl RbDataFrame {
                 .map(|s| match s.dtype() {
                     DataType::Object(_, _) => {
                         let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
-                        obj.unwrap().to_object()
+                        obj.unwrap().to_value()
                     }
                     _ => Wrap(s.get(idx).unwrap()).into_value(),
                 }),
@@ -37,7 +39,7 @@ impl RbDataFrame {
                     .map(|s| match s.dtype() {
                         DataType::Object(_, _) => {
                             let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
-                            obj.unwrap().to_object()
+                            obj.unwrap().to_value()
                         }
                         _ => Wrap(s.get(idx).unwrap()).into_value(),
                     }),
@@ -45,4 +47,9 @@ impl RbDataFrame {
         }))
         .as_value()
     }
+
+    pub fn __arrow_c_stream__(&self) -> RbResult<Value> {
+        self.df.borrow_mut().align_chunks();
+        dataframe_to_stream(&self.df.borrow())
+    }
 }

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

@@ -1,6 +1,5 @@
 use either::Either;
 use magnus::{prelude::*, typed_data::Obj, IntoValue, RArray, Value};
-use polars::frame::NullStrategy;
 use polars::prelude::pivot::{pivot, pivot_stable};
 use polars::prelude::*;
 
@@ -10,14 +9,14 @@ use crate::map::dataframe::{
     apply_lambda_with_utf8_out_type,
 };
 use crate::prelude::strings_to_pl_smallstr;
-use crate::series::{to_rbseries_collection, to_series_collection};
+use crate::series::{to_rbseries, to_series};
 use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
 
 impl RbDataFrame {
     pub fn init(columns: RArray) -> RbResult<Self> {
         let mut cols = Vec::new();
         for i in columns.into_iter() {
-            cols.push(<&RbSeries>::try_convert(i)?.series.borrow().clone());
+            cols.push(<&RbSeries>::try_convert(i)?.series.borrow().clone().into());
         }
         let df = DataFrame::new(cols).map_err(RbPolarsErr::from)?;
         Ok(RbDataFrame::new(df))
@@ -128,7 +127,7 @@ impl RbDataFrame {
 
     pub fn get_columns(&self) -> RArray {
         let cols = self.df.borrow().get_columns().to_vec();
-        to_rbseries_collection(cols)
+        to_rbseries(cols)
     }
 
     pub fn columns(&self) -> Vec<String> {
@@ -158,7 +157,7 @@ impl RbDataFrame {
     }
 
     pub fn n_chunks(&self) -> usize {
-        self.df.borrow().n_chunks()
+        self.df.borrow().first_col_n_chunks()
     }
 
     pub fn shape(&self) -> (usize, usize) {
@@ -174,7 +173,8 @@ impl RbDataFrame {
     }
 
     pub fn hstack(&self, columns: RArray) -> RbResult<Self> {
-        let columns = to_series_collection(columns)?;
+        let columns = to_series(columns)?;
+        let columns = columns.into_iter().map(Into::into).collect::<Vec<_>>();
         let df = self
             .df
             .borrow()
@@ -184,7 +184,8 @@ impl RbDataFrame {
     }
 
     pub fn hstack_mut(&self, columns: RArray) -> RbResult<()> {
-        let columns = to_series_collection(columns)?;
+        let columns = to_series(columns)?;
+        let columns = columns.into_iter().map(Into::into).collect::<Vec<_>>();
         self.df
             .borrow_mut()
             .hstack_mut(&columns)
@@ -223,6 +224,7 @@ impl RbDataFrame {
             .borrow_mut()
             .drop_in_place(&name)
             .map_err(RbPolarsErr::from)?;
+        let s = s.take_materialized_series();
         Ok(RbSeries::new(s))
     }
 
@@ -230,7 +232,7 @@ impl RbDataFrame {
         self.df
             .borrow()
             .select_at_idx(idx)
-            .map(|s| RbSeries::new(s.clone()))
+            .map(|s| RbSeries::new(s.as_materialized_series().clone()))
     }
 
     pub fn get_column_index(&self, name: String) -> Option<usize> {
@@ -238,11 +240,13 @@ impl RbDataFrame {
     }
 
     pub fn get_column(&self, name: String) -> RbResult<RbSeries> {
-        self.df
+        let series = self
+            .df
             .borrow()
             .column(&name)
-            .map(|s| RbSeries::new(s.clone()))
-            .map_err(RbPolarsErr::from)
+            .map(|s| RbSeries::new(s.as_materialized_series().clone()))
+            .map_err(RbPolarsErr::from)?;
+        Ok(series)
     }
 
     pub fn select(&self, selection: Vec<String>) -> RbResult<Self> {
@@ -405,52 +409,6 @@ impl RbDataFrame {
         self.df.borrow().clone().lazy().into()
     }
 
-    pub fn max_horizontal(&self) -> RbResult<Option<RbSeries>> {
-        let s = self
-            .df
-            .borrow()
-            .max_horizontal()
-            .map_err(RbPolarsErr::from)?;
-        Ok(s.map(|s| s.into()))
-    }
-
-    pub fn min_horizontal(&self) -> RbResult<Option<RbSeries>> {
-        let s = self
-            .df
-            .borrow()
-            .min_horizontal()
-            .map_err(RbPolarsErr::from)?;
-        Ok(s.map(|s| s.into()))
-    }
-
-    pub fn sum_horizontal(&self, ignore_nulls: bool) -> RbResult<Option<RbSeries>> {
-        let null_strategy = if ignore_nulls {
-            NullStrategy::Ignore
-        } else {
-            NullStrategy::Propagate
-        };
-        let s = self
-            .df
-            .borrow()
-            .sum_horizontal(null_strategy)
-            .map_err(RbPolarsErr::from)?;
-        Ok(s.map(|s| s.into()))
-    }
-
-    pub fn mean_horizontal(&self, ignore_nulls: bool) -> RbResult<Option<RbSeries>> {
-        let null_strategy = if ignore_nulls {
-            NullStrategy::Ignore
-        } else {
-            NullStrategy::Propagate
-        };
-        let s = self
-            .df
-            .borrow()
-            .mean_horizontal(null_strategy)
-            .map_err(RbPolarsErr::from)?;
-        Ok(s.map(|s| s.into()))
-    }
-
     pub fn to_dummies(
         &self,
         columns: Option<Vec<String>>,