polars-df 0.20.0 → 0.21.1

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

@@ -2,27 +2,37 @@ pub mod dataframe;
 pub mod lazy;
 pub mod series;
 
-use magnus::{prelude::*, RHash, Value};
+use magnus::{RHash, Value, prelude::*};
 use polars::chunked_array::builder::get_list_builder;
 use polars::prelude::*;
-use polars_core::utils::CustomIterTools;
 use polars_core::POOL;
+use polars_core::utils::CustomIterTools;
 use rayon::prelude::*;
 
-use crate::{ObjectValue, RbPolarsErr, RbResult, RbSeries, Wrap};
+use crate::{ObjectValue, RbPolarsErr, RbResult, RbSeries, RbValueError, Wrap};
+
+pub trait RbPolarsNumericType: PolarsNumericType {}
 
-pub trait RbArrowPrimitiveType: PolarsNumericType {}
+impl RbPolarsNumericType for UInt8Type {}
+impl RbPolarsNumericType for UInt16Type {}
+impl RbPolarsNumericType for UInt32Type {}
+impl RbPolarsNumericType for UInt64Type {}
+impl RbPolarsNumericType for Int8Type {}
+impl RbPolarsNumericType for Int16Type {}
+impl RbPolarsNumericType for Int32Type {}
+impl RbPolarsNumericType for Int64Type {}
+impl RbPolarsNumericType for Float32Type {}
+impl RbPolarsNumericType for Float64Type {}
 
-impl RbArrowPrimitiveType for UInt8Type {}
-impl RbArrowPrimitiveType for UInt16Type {}
-impl RbArrowPrimitiveType for UInt32Type {}
-impl RbArrowPrimitiveType for UInt64Type {}
-impl RbArrowPrimitiveType for Int8Type {}
-impl RbArrowPrimitiveType for Int16Type {}
-impl RbArrowPrimitiveType for Int32Type {}
-impl RbArrowPrimitiveType for Int64Type {}
-impl RbArrowPrimitiveType for Float32Type {}
-impl RbArrowPrimitiveType for Float64Type {}
+pub(super) fn check_nested_object(dt: &DataType) -> RbResult<()> {
+    if dt.contains_objects() {
+        Err(RbValueError::new_err(
+            "nested objects are not allowed\n\nSet `return_dtype: Polars::Object` to use Ruby's native nesting.",
+        ))
+    } else {
+        Ok(())
+    }
+}
 
 fn iterator_to_struct(
     it: impl Iterator<Item = Option<Value>>,
@@ -37,7 +47,7 @@ fn iterator_to_struct(
         _ => {
             return Err(crate::exceptions::ComputeError::new_err(format!(
                 "expected struct got {first_value:?}",
-            )))
+            )));
         }
     };
 
@@ -70,9 +80,11 @@ fn iterator_to_struct(
             Some(dict) => {
                 let dict = RHash::try_convert(dict)?;
                 if dict.len() != struct_width {
-                    return Err(crate::exceptions::ComputeError::new_err(
-                        format!("Cannot create struct type.\n> The struct dtype expects {} fields, but it got a dict with {} fields.", struct_width, dict.len())
-                    ));
+                    return Err(crate::exceptions::ComputeError::new_err(format!(
+                        "Cannot create struct type.\n> The struct dtype expects {} fields, but it got a dict with {} fields.",
+                        struct_width,
+                        dict.len()
+                    )));
                 }
                 // we ignore the keys of the rest of the dicts
                 // the first item determines the output name
@@ -109,7 +121,7 @@ fn iterator_to_primitive<T>(
     capacity: usize,
 ) -> ChunkedArray<T>
 where
-    T: RbArrowPrimitiveType,
+    T: RbPolarsNumericType,
 {
     // safety: we know the iterators len
     let mut ca: ChunkedArray<T> = unsafe {

data/ext/polars/src/map/series.rs

@@ -1,4 +1,4 @@
-use magnus::{class, prelude::*, typed_data::Obj, IntoValue, TryConvert, Value};
+use magnus::{IntoValue, TryConvert, Value, class, prelude::*, typed_data::Obj};
 use polars::prelude::*;
 
 use super::*;
@@ -85,7 +85,7 @@ pub trait ApplyLambda<'a> {
         first_value: Option<D::Native>,
     ) -> RbResult<ChunkedArray<D>>
     where
-        D: RbArrowPrimitiveType,
+        D: RbPolarsNumericType,
         D::Native: IntoValue + TryConvert;
 
     /// Apply a lambda with a boolean output type
@@ -143,7 +143,7 @@ where
 {
     match call_lambda(lambda, in_val) {
         Ok(out) => S::try_convert(out),
-        Err(e) => panic!("ruby function failed {}", e),
+        Err(e) => panic!("ruby function failed {e}"),
     }
 }
 
@@ -219,7 +219,7 @@ impl<'a> ApplyLambda<'a> for BooleanChunked {
         first_value: Option<D::Native>,
     ) -> RbResult<ChunkedArray<D>>
     where
-        D: RbArrowPrimitiveType,
+        D: RbPolarsNumericType,
         D::Native: IntoValue + TryConvert,
     {
         let skip = usize::from(first_value.is_some());
@@ -438,7 +438,7 @@ impl<'a> ApplyLambda<'a> for BooleanChunked {
 
 impl<'a, T> ApplyLambda<'a> for ChunkedArray<T>
 where
-    T: RbArrowPrimitiveType + PolarsNumericType,
+    T: RbPolarsNumericType + PolarsNumericType,
     T::Native: IntoValue + TryConvert,
     ChunkedArray<T>: IntoSeries,
 {
@@ -503,7 +503,7 @@ where
         first_value: Option<D::Native>,
     ) -> RbResult<ChunkedArray<D>>
     where
-        D: RbArrowPrimitiveType,
+        D: RbPolarsNumericType,
         D::Native: IntoValue + TryConvert,
     {
         let skip = usize::from(first_value.is_some());
@@ -782,7 +782,7 @@ impl<'a> ApplyLambda<'a> for StringChunked {
         first_value: Option<D::Native>,
     ) -> RbResult<ChunkedArray<D>>
     where
-        D: RbArrowPrimitiveType,
+        D: RbPolarsNumericType,
         D::Native: IntoValue + TryConvert,
     {
         let skip = usize::from(first_value.is_some());
@@ -999,7 +999,7 @@ impl<'a> ApplyLambda<'a> for StringChunked {
     }
 }
 
-fn iter_struct(ca: &StructChunked) -> impl Iterator<Item = AnyValue> {
+fn iter_struct(ca: &StructChunked) -> impl Iterator<Item = AnyValue<'_>> {
     (0..ca.len()).map(|i| unsafe { ca.get_any_value_unchecked(i) })
 }
 
@@ -1047,7 +1047,7 @@ impl<'a> ApplyLambda<'a> for StructChunked {
         first_value: Option<D::Native>,
     ) -> RbResult<ChunkedArray<D>>
     where
-        D: RbArrowPrimitiveType,
+        D: RbPolarsNumericType,
         D::Native: IntoValue + TryConvert,
     {
         let skip = usize::from(first_value.is_some());

data/ext/polars/src/on_startup.rs

@@ -9,8 +9,8 @@ use polars_core::chunked_array::object::registry;
 use polars_core::chunked_array::object::registry::AnonymousObjectBuilder;
 use polars_core::prelude::AnyValue;
 
-use crate::prelude::ObjectValue;
 use crate::Wrap;
+use crate::prelude::ObjectValue;
 
 static POLARS_REGISTRY_INIT_LOCK: OnceLock<()> = OnceLock::new();
 
@@ -39,6 +39,9 @@ pub(crate) fn register_startup_deps() {
             object_converter,
             rbobject_converter,
             physical_dtype,
-        )
+        );
+        // TODO
+        // Register warning function for `polars_warn!`.
+        // polars_error::set_warning_function(warning_function);
     });
 }

data/ext/polars/src/rb_modules.rs

@@ -1,4 +1,4 @@
-use magnus::{value::Lazy, ExceptionClass, Module, RClass, RModule, Ruby};
+use magnus::{ExceptionClass, Module, RClass, RModule, Ruby, value::Lazy};
 
 static POLARS: Lazy<RModule> = Lazy::new(|ruby| ruby.class_object().const_get("Polars").unwrap());
 

data/ext/polars/src/series/aggregation.rs

@@ -3,6 +3,10 @@ use crate::prelude::*;
 use crate::{RbResult, RbSeries};
 use magnus::{IntoValue, Value};
 
+fn scalar_to_rb(scalar: RbResult<Scalar>) -> RbResult<Value> {
+    Ok(Wrap(scalar?.as_any_value()).into_value())
+}
+
 impl RbSeries {
     pub fn any(&self, ignore_nulls: bool) -> RbResult<Option<bool>> {
         let binding = self.series.borrow();
@@ -118,4 +122,44 @@ impl RbSeries {
         )
         .into_value())
     }
+
+    pub fn first(&self) -> RbResult<Value> {
+        scalar_to_rb(Ok(self.series.borrow().first()))
+    }
+
+    pub fn last(&self) -> RbResult<Value> {
+        scalar_to_rb(Ok(self.series.borrow().last()))
+    }
+
+    pub fn approx_n_unique(&self) -> RbResult<IdxSize> {
+        Ok(self
+            .series
+            .borrow()
+            .approx_n_unique()
+            .map_err(RbPolarsErr::from)?)
+    }
+
+    pub fn bitwise_and(&self) -> RbResult<Value> {
+        scalar_to_rb(Ok(self
+            .series
+            .borrow()
+            .and_reduce()
+            .map_err(RbPolarsErr::from)?))
+    }
+
+    pub fn bitwise_or(&self) -> RbResult<Value> {
+        scalar_to_rb(Ok(self
+            .series
+            .borrow()
+            .or_reduce()
+            .map_err(RbPolarsErr::from)?))
+    }
+
+    pub fn bitwise_xor(&self) -> RbResult<Value> {
+        scalar_to_rb(Ok(self
+            .series
+            .borrow()
+            .xor_reduce()
+            .map_err(RbPolarsErr::from)?))
+    }
 }

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

@@ -1,8 +1,8 @@
-use magnus::{prelude::*, RArray, RString};
+use magnus::{RArray, RString, prelude::*};
 use polars_core::prelude::*;
 
 use crate::any_value::rb_object_to_any_value;
-use crate::conversion::{slice_extract_wrapped, vec_extract_wrapped, Wrap};
+use crate::conversion::{Wrap, slice_extract_wrapped, vec_extract_wrapped};
 use crate::prelude::ObjectValue;
 use crate::series::to_series;
 use crate::{RbPolarsErr, RbResult, RbSeries, RbTypeError, RbValueError};
@@ -122,13 +122,17 @@ impl RbSeries {
             .into_iter()
             .map(|v| rb_object_to_any_value(v, strict))
             .collect::<RbResult<Vec<AnyValue>>>()?;
-        let s =
-            Series::from_any_values_and_dtype(name.into(), any_values.as_slice(), &dtype.0, strict)
-                .map_err(|e| {
-                    RbTypeError::new_err(format!(
+        let s = Series::from_any_values_and_dtype(
+            name.into(),
+            any_values.as_slice(),
+            &dtype.0,
+            strict,
+        )
+        .map_err(|e| {
+            RbTypeError::new_err(format!(
                 "{e}\n\nHint: Try setting `strict: false` to allow passing data with mixed types."
             ))
-                })?;
+        })?;
         Ok(s.into())
     }
 

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

@@ -1,8 +1,8 @@
-use magnus::{value::qnil, IntoValue, RArray, Value};
+use magnus::{IntoValue, RArray, Value, value::qnil};
 use polars_core::prelude::*;
 
-use crate::prelude::*;
 use crate::RbSeries;
+use crate::prelude::*;
 
 impl RbSeries {
     /// Convert this Series to a Ruby array.
@@ -25,7 +25,9 @@ impl RbSeries {
                 DataType::Float32 => RArray::from_iter(series.f32().unwrap()).into_value(),
                 DataType::Float64 => RArray::from_iter(series.f64().unwrap()).into_value(),
                 DataType::Categorical(_, _) | DataType::Enum(_, _) => {
-                    RArray::from_iter(series.categorical().unwrap().iter_str()).into_value()
+                    with_match_categorical_physical_type!(series.dtype().cat_physical().unwrap(), |$C| {
+                        RArray::from_iter(series.cat::<$C>().unwrap().iter_str()).into_value()
+                    })
                 }
                 DataType::Object(_) => {
                     let v = RArray::with_capacity(series.len());
@@ -132,7 +134,7 @@ impl RbSeries {
                     unreachable!()
                 }
             };
-            rblist
+            rblist.into_value()
         }
 
         to_a_recursive(series)

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

@@ -1,10 +1,11 @@
-use magnus::{exception, Error, IntoValue, Value};
+use magnus::{Error, IntoValue, RArray, Value, exception, value::ReprValue};
 use polars::prelude::*;
 use polars::series::IsSorted;
+use polars_core::utils::flatten::flatten_series;
 
-use crate::apply_method_all_arrow_series2;
 use crate::conversion::*;
-use crate::map::series::{call_lambda_and_extract, ApplyLambda};
+use crate::exceptions::RbIndexError;
+use crate::rb_modules;
 use crate::{RbDataFrame, RbPolarsErr, RbResult, RbSeries};
 
 impl RbSeries {
@@ -42,21 +43,15 @@ impl RbSeries {
     }
 
     pub fn cat_uses_lexical_ordering(&self) -> RbResult<bool> {
-        let binding = self.series.borrow();
-        let ca = binding.categorical().map_err(RbPolarsErr::from)?;
-        Ok(ca.uses_lexical_ordering())
+        Ok(true)
     }
 
     pub fn cat_is_local(&self) -> RbResult<bool> {
-        let binding = self.series.borrow();
-        let ca = binding.categorical().map_err(RbPolarsErr::from)?;
-        Ok(ca.get_rev_map().is_local())
+        Ok(false)
     }
 
     pub fn cat_to_local(&self) -> RbResult<Self> {
-        let binding = self.series.borrow();
-        let ca = binding.categorical().map_err(RbPolarsErr::from)?;
-        Ok(ca.to_local().into_series().into())
+        Ok(self.clone())
     }
 
     pub fn estimated_size(&self) -> usize {
@@ -75,7 +70,7 @@ impl RbSeries {
             if val == v_trunc {
                 val
             } else {
-                format!("{}…", v_trunc)
+                format!("{v_trunc}…")
             }
         } else {
             val
@@ -92,8 +87,39 @@ impl RbSeries {
         }
     }
 
-    pub fn get_idx(&self, idx: usize) -> RbResult<Value> {
-        Ok(Wrap(self.series.borrow().get(idx).map_err(RbPolarsErr::from)?).into_value())
+    pub fn get_index(&self, index: usize) -> RbResult<Value> {
+        let binding = self.series.borrow();
+        let av = match binding.get(index) {
+            Ok(v) => v,
+            Err(PolarsError::OutOfBounds(err)) => {
+                return Err(RbIndexError::new_err(err.to_string()));
+            }
+            Err(e) => return Err(RbPolarsErr::from(e).into()),
+        };
+
+        match av {
+            AnyValue::List(s) | AnyValue::Array(s, _) => {
+                let rbseries = RbSeries::new(s);
+                rb_modules::utils().funcall("wrap_s", (rbseries,))
+            }
+            _ => Ok(Wrap(av).into_value()),
+        }
+    }
+
+    pub fn get_index_signed(&self, index: isize) -> RbResult<Value> {
+        let index = if index < 0 {
+            match self.len().checked_sub(index.unsigned_abs()) {
+                Some(v) => v,
+                None => {
+                    return Err(RbIndexError::new_err(
+                        polars_err!(oob = index, self.len()).to_string(),
+                    ));
+                }
+            }
+        } else {
+            usize::try_from(index).unwrap()
+        };
+        self.get_index(index)
     }
 
     pub fn bitand(&self, other: &RbSeries) -> RbResult<Self> {
@@ -287,7 +313,7 @@ impl RbSeries {
         }
     }
 
-    pub fn not(&self) -> RbResult<Self> {
+    pub fn not_(&self) -> RbResult<Self> {
         let binding = self.series.borrow();
         let bool = binding.bool().map_err(RbPolarsErr::from)?;
         Ok((!bool).into_series().into())
@@ -305,198 +331,6 @@ impl RbSeries {
         RbSeries::new(self.series.borrow().clone())
     }
 
-    pub fn apply_lambda(
-        &self,
-        lambda: Value,
-        output_type: Option<Wrap<DataType>>,
-        skip_nulls: bool,
-    ) -> RbResult<Self> {
-        let series = &self.series.borrow();
-
-        let output_type = output_type.map(|dt| dt.0);
-
-        macro_rules! dispatch_apply {
-            ($self:expr, $method:ident, $($args:expr),*) => {
-                if matches!($self.dtype(), DataType::Object(_)) {
-                    // let ca = $self.0.unpack::<ObjectType<ObjectValue>>().unwrap();
-                    // ca.$method($($args),*)
-                    todo!()
-                } else {
-                    apply_method_all_arrow_series2!(
-                        $self,
-                        $method,
-                        $($args),*
-                    )
-                }
-
-            }
-
-        }
-
-        if matches!(
-            series.dtype(),
-            DataType::Datetime(_, _)
-                | DataType::Date
-                | DataType::Duration(_)
-                | DataType::Categorical(_, _)
-                | DataType::Time
-        ) || !skip_nulls
-        {
-            let mut avs = Vec::with_capacity(series.len());
-            let iter = series.iter().map(|av| {
-                let input = Wrap(av);
-                call_lambda_and_extract::<_, Wrap<AnyValue>>(lambda, input)
-                    .unwrap()
-                    .0
-            });
-            avs.extend(iter);
-            return Ok(Series::new(self.name().into(), &avs).into());
-        }
-
-        let out = match output_type {
-            Some(DataType::Int8) => {
-                let ca: Int8Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Int16) => {
-                let ca: Int16Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Int32) => {
-                let ca: Int32Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Int64) => {
-                let ca: Int64Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::UInt8) => {
-                let ca: UInt8Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::UInt16) => {
-                let ca: UInt16Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::UInt32) => {
-                let ca: UInt32Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::UInt64) => {
-                let ca: UInt64Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Float32) => {
-                let ca: Float32Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Float64) => {
-                let ca: Float64Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_series()
-            }
-            Some(DataType::Boolean) => {
-                let ca: BooleanChunked =
-                    dispatch_apply!(series, apply_lambda_with_bool_out_type, lambda, 0, None)?;
-                ca.into_series()
-            }
-            Some(DataType::Date) => {
-                let ca: Int32Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_date().into_series()
-            }
-            Some(DataType::Datetime(tu, tz)) => {
-                let ca: Int64Chunked = dispatch_apply!(
-                    series,
-                    apply_lambda_with_primitive_out_type,
-                    lambda,
-                    0,
-                    None
-                )?;
-                ca.into_datetime(tu, tz).into_series()
-            }
-            Some(DataType::String) => {
-                let ca = dispatch_apply!(series, apply_lambda_with_utf8_out_type, lambda, 0, None)?;
-
-                ca.into_series()
-            }
-            Some(DataType::Object(_)) => {
-                let ca =
-                    dispatch_apply!(series, apply_lambda_with_object_out_type, lambda, 0, None)?;
-                ca.into_series()
-            }
-            None => return dispatch_apply!(series, apply_lambda_unknown, lambda),
-
-            _ => return dispatch_apply!(series, apply_lambda_unknown, lambda),
-        };
-
-        Ok(RbSeries::new(out))
-    }
-
     pub fn zip_with(&self, mask: &RbSeries, other: &RbSeries) -> RbResult<Self> {
         let binding = mask.series.borrow();
         let mask = binding.bool().map_err(RbPolarsErr::from)?;
@@ -508,11 +342,16 @@ impl RbSeries {
         Ok(RbSeries::new(s))
     }
 
-    pub fn to_dummies(&self, sep: Option<String>, drop_first: bool) -> RbResult<RbDataFrame> {
+    pub fn to_dummies(
+        &self,
+        sep: Option<String>,
+        drop_first: bool,
+        drop_nulls: bool,
+    ) -> RbResult<RbDataFrame> {
         let df = self
             .series
             .borrow()
-            .to_dummies(sep.as_deref(), drop_first)
+            .to_dummies(sep.as_deref(), drop_first, drop_nulls)
             .map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }
@@ -565,6 +404,32 @@ impl RbSeries {
         Ok(out.into())
     }
 
+    pub fn get_chunks(&self) -> RbResult<RArray> {
+        flatten_series(&self.series.borrow())
+            .into_iter()
+            .map(|s| rb_modules::utils().funcall::<_, _, Value>("wrap_s", (Self::new(s),)))
+            .collect()
+    }
+
+    pub fn is_sorted(&self, descending: bool, nulls_last: bool) -> RbResult<bool> {
+        let options = SortOptions {
+            descending,
+            nulls_last,
+            multithreaded: true,
+            maintain_order: false,
+            limit: None,
+        };
+        Ok(self
+            .series
+            .borrow()
+            .is_sorted(options)
+            .map_err(RbPolarsErr::from)?)
+    }
+
+    pub fn clear(&self) -> Self {
+        self.series.borrow().clear().into()
+    }
+
     pub fn time_unit(&self) -> Option<String> {
         if let DataType::Datetime(tu, _) | DataType::Duration(tu) = self.series.borrow().dtype() {
             Some(

data/ext/polars/src/series/import.rs

@@ -1,13 +1,13 @@
 use arrow::array::Array;
 use arrow::ffi::{ArrowArrayStream, ArrowArrayStreamReader};
-use magnus::prelude::*;
 use magnus::Value;
+use magnus::prelude::*;
 use polars::prelude::*;
 
 use super::RbSeries;
 
-use crate::exceptions::RbValueError;
 use crate::RbResult;
+use crate::exceptions::RbValueError;
 
 /// Import `arrow_c_stream` across Ruby boundary.
 fn call_arrow_c_stream(ob: Value) -> RbResult<Value> {