polars-df 0.7.0 → 0.9.0

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

@@ -187,7 +187,7 @@ pub fn apply_lambda_with_utf8_out_type(
     lambda: Value,
     init_null_count: usize,
     first_value: Option<&str>,
-) -> Utf8Chunked {
+) -> StringChunked {
     let skip = usize::from(first_value.is_some());
     if init_null_count == df.height() {
         ChunkedArray::full_null("apply", df.height())

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

@@ -198,11 +198,11 @@ fn iterator_to_utf8(
     first_value: Option<&str>,
     name: &str,
     capacity: usize,
-) -> Utf8Chunked {
+) -> StringChunked {
     let first_value = first_value.map(|v| v.to_string());
 
     // safety: we know the iterators len
-    let mut ca: Utf8Chunked = unsafe {
+    let mut ca: StringChunked = unsafe {
         if init_null_count > 0 {
             (0..init_null_count)
                 .map(|_| None)

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

@@ -101,7 +101,7 @@ pub trait ApplyLambda<'a> {
         lambda: Value,
         init_null_count: usize,
         first_value: Option<&str>,
-    ) -> RbResult<Utf8Chunked>;
+    ) -> RbResult<StringChunked>;
 
     /// Apply a lambda with list output type
     fn apply_lambda_with_list_out_type(
@@ -285,7 +285,7 @@ impl<'a> ApplyLambda<'a> for BooleanChunked {
         lambda: Value,
         init_null_count: usize,
         first_value: Option<&str>,
-    ) -> RbResult<Utf8Chunked> {
+    ) -> RbResult<StringChunked> {
         let skip = usize::from(first_value.is_some());
         if init_null_count == self.len() {
             Ok(ChunkedArray::full_null(self.name(), self.len()))
@@ -562,7 +562,7 @@ where
         lambda: Value,
         init_null_count: usize,
         first_value: Option<&str>,
-    ) -> RbResult<Utf8Chunked> {
+    ) -> RbResult<StringChunked> {
         let skip = usize::from(first_value.is_some());
         if init_null_count == self.len() {
             Ok(ChunkedArray::full_null(self.name(), self.len()))
@@ -705,7 +705,7 @@ where
     }
 }
 
-impl<'a> ApplyLambda<'a> for Utf8Chunked {
+impl<'a> ApplyLambda<'a> for StringChunked {
     fn apply_lambda_unknown(&'a self, lambda: Value) -> RbResult<RbSeries> {
         let mut null_count = 0;
         for opt_v in self.into_iter() {
@@ -834,7 +834,7 @@ impl<'a> ApplyLambda<'a> for Utf8Chunked {
         lambda: Value,
         init_null_count: usize,
         first_value: Option<&str>,
-    ) -> RbResult<Utf8Chunked> {
+    ) -> RbResult<StringChunked> {
         let skip = usize::from(first_value.is_some());
         if init_null_count == self.len() {
             Ok(ChunkedArray::full_null(self.name(), self.len()))
@@ -1079,7 +1079,7 @@ impl<'a> ApplyLambda<'a> for StructChunked {
         lambda: Value,
         init_null_count: usize,
         first_value: Option<&str>,
-    ) -> RbResult<Utf8Chunked> {
+    ) -> RbResult<StringChunked> {
         let names = self.fields().iter().map(|s| s.name()).collect::<Vec<_>>();
 
         let skip = usize::from(first_value.is_some());

data/ext/polars/src/object.rs

@@ -1,31 +1 @@
-use std::any::Any;
-use std::sync::Arc;
-
-use magnus::IntoValue;
-use polars_core::chunked_array::object::builder::ObjectChunkedBuilder;
-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;
-
 pub(crate) const OBJECT_NAME: &str = "object";
-
-pub(crate) fn register_object_builder() {
-    if !registry::is_object_builder_registered() {
-        let object_builder = Box::new(|name: &str, capacity: usize| {
-            Box::new(ObjectChunkedBuilder::<ObjectValue>::new(name, capacity))
-                as Box<dyn AnonymousObjectBuilder>
-        });
-
-        let object_converter = Arc::new(|av: AnyValue| {
-            let object = ObjectValue {
-                inner: Wrap(av).into_value().into(),
-            };
-            Box::new(object) as Box<dyn Any>
-        });
-
-        registry::register_object_builder(object_builder, object_converter)
-    }
-}

data/ext/polars/src/on_startup.rs

@@ -0,0 +1,32 @@
+use std::any::Any;
+use std::sync::Arc;
+
+use magnus::IntoValue;
+use polars::prelude::*;
+use polars_core::chunked_array::object::builder::ObjectChunkedBuilder;
+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;
+
+pub(crate) fn register_object_builder() {
+    if !registry::is_object_builder_registered() {
+        let object_builder = Box::new(|name: &str, capacity: usize| {
+            Box::new(ObjectChunkedBuilder::<ObjectValue>::new(name, capacity))
+                as Box<dyn AnonymousObjectBuilder>
+        });
+
+        let object_converter = Arc::new(|av: AnyValue| {
+            let object = ObjectValue {
+                inner: Wrap(av).into_value().into(),
+            };
+            Box::new(object) as Box<dyn Any>
+        });
+
+        let object_size = std::mem::size_of::<ObjectValue>();
+        let physical_dtype = ArrowDataType::FixedSizeBinary(object_size);
+        registry::register_object_builder(object_builder, object_converter, physical_dtype)
+    }
+}

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

@@ -4,6 +4,26 @@ use crate::{RbResult, RbSeries, RbValueError};
 use magnus::{IntoValue, Value};
 
 impl RbSeries {
+    pub fn any(&self, ignore_nulls: bool) -> RbResult<Option<bool>> {
+        let binding = self.series.borrow();
+        let s = binding.bool().map_err(RbPolarsErr::from)?;
+        Ok(if ignore_nulls {
+            Some(s.any())
+        } else {
+            s.any_kleene()
+        })
+    }
+
+    pub fn all(&self, ignore_nulls: bool) -> RbResult<Option<bool>> {
+        let binding = self.series.borrow();
+        let s = binding.bool().map_err(RbPolarsErr::from)?;
+        Ok(if ignore_nulls {
+            Some(s.all())
+        } else {
+            s.all_kleene()
+        })
+    }
+
     pub fn arg_max(&self) -> Option<usize> {
         self.series.borrow().arg_max()
     }
@@ -17,6 +37,7 @@ impl RbSeries {
             self.series
                 .borrow()
                 .max_as_series()
+                .map_err(RbPolarsErr::from)?
                 .get(0)
                 .map_err(RbPolarsErr::from)?,
         )
@@ -48,6 +69,7 @@ impl RbSeries {
             self.series
                 .borrow()
                 .min_as_series()
+                .map_err(RbPolarsErr::from)?
                 .get(0)
                 .map_err(RbPolarsErr::from)?,
         )
@@ -75,6 +97,7 @@ impl RbSeries {
             self.series
                 .borrow()
                 .sum_as_series()
+                .map_err(RbPolarsErr::from)?
                 .get(0)
                 .map_err(RbPolarsErr::from)?,
         )

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

@@ -106,7 +106,7 @@ impl RbSeries {
         Ok(s.into())
     }
 
-    pub fn new_str(name: String, val: Wrap<Utf8Chunked>, _strict: bool) -> Self {
+    pub fn new_str(name: String, val: Wrap<StringChunked>, _strict: bool) -> Self {
         let mut s = val.0.into_series();
         s.rename(&name);
         RbSeries::new(s)

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

@@ -9,8 +9,8 @@ impl RbSeries {
     pub fn to_numo(&self) -> RbResult<Value> {
         let s = &self.series.borrow();
         match s.dtype() {
-            DataType::Utf8 => {
-                let ca = s.utf8().unwrap();
+            DataType::String => {
+                let ca = s.str().unwrap();
 
                 // TODO make more efficient
                 let np_arr = RArray::from_iter(ca);

data/ext/polars/src/{series.rs → series/mod.rs}

@@ -3,7 +3,7 @@ mod arithmetic;
 mod comparison;
 mod construction;
 mod export;
-mod set_at_idx;
+mod scatter;
 
 use magnus::{exception, prelude::*, value::qnil, Error, IntoValue, RArray, Value};
 use polars::prelude::*;
@@ -74,13 +74,31 @@ 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())
+    }
+
+    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())
+    }
+
+    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())
+    }
+
     pub fn estimated_size(&self) -> usize {
         self.series.borrow().estimated_size()
     }
 
     pub fn get_fmt(&self, index: usize, str_lengths: usize) -> String {
         let val = format!("{}", self.series.borrow().get(index).unwrap());
-        if let DataType::Utf8 | DataType::Categorical(_) = self.series.borrow().dtype() {
+        if let DataType::String | DataType::Categorical(_, _) = self.series.borrow().dtype() {
             let v_trunc = &val[..val
                 .char_indices()
                 .take(str_lengths)
@@ -90,7 +108,7 @@ impl RbSeries {
             if val == v_trunc {
                 val
             } else {
-                format!("{}...", v_trunc)
+                format!("{}…", v_trunc)
             }
         } else {
             val
@@ -215,8 +233,8 @@ impl RbSeries {
         }
     }
 
-    pub fn sort(&self, reverse: bool) -> Self {
-        (self.series.borrow_mut().sort(reverse)).into()
+    pub fn sort(&self, descending: bool, nulls_last: bool) -> Self {
+        (self.series.borrow_mut().sort(descending, nulls_last)).into()
     }
 
     pub fn value_counts(&self, sorted: bool) -> RbResult<RbDataFrame> {
@@ -273,15 +291,13 @@ impl RbSeries {
         Ok(s.into())
     }
 
-    pub fn series_equal(&self, other: &RbSeries, null_equal: bool, strict: bool) -> bool {
+    pub fn equals(&self, other: &RbSeries, null_equal: bool, strict: bool) -> bool {
         if strict {
             self.series.borrow().eq(&other.series.borrow())
         } else if null_equal {
-            self.series
-                .borrow()
-                .series_equal_missing(&other.series.borrow())
+            self.series.borrow().equals_missing(&other.series.borrow())
         } else {
-            self.series.borrow().series_equal(&other.series.borrow())
+            self.series.borrow().equals(&other.series.borrow())
         }
     }
 
@@ -315,10 +331,10 @@ impl RbSeries {
                 DataType::Int64 => RArray::from_iter(series.i64().unwrap()).into_value(),
                 DataType::Float32 => RArray::from_iter(series.f32().unwrap()).into_value(),
                 DataType::Float64 => RArray::from_iter(series.f64().unwrap()).into_value(),
-                DataType::Categorical(_) => {
+                DataType::Categorical(_, _) | DataType::Enum(_, _) => {
                     RArray::from_iter(series.categorical().unwrap().iter_str()).into_value()
                 }
-                DataType::Object(_) => {
+                DataType::Object(_, _) => {
                     let v = RArray::with_capacity(series.len());
                     for i in 0..series.len() {
                         let obj: Option<&ObjectValue> = series.get_object(i).map(|any| any.into());
@@ -377,8 +393,8 @@ impl RbSeries {
                     let ca = series.decimal().unwrap();
                     return Wrap(ca).into_value();
                 }
-                DataType::Utf8 => {
-                    let ca = series.utf8().unwrap();
+                DataType::String => {
+                    let ca = series.str().unwrap();
                     return Wrap(ca).into_value();
                 }
                 DataType::Struct(_) => {
@@ -417,7 +433,10 @@ impl RbSeries {
                     RArray::from_iter(NullIter { iter, n }).into_value()
                 }
                 DataType::Unknown => {
-                    panic!("to_a not implemented for null/unknown")
+                    panic!("to_a not implemented for unknown")
+                }
+                DataType::BinaryOffset => {
+                    unreachable!()
                 }
             };
             rblist
@@ -442,7 +461,7 @@ impl RbSeries {
 
         macro_rules! dispatch_apply {
             ($self:expr, $method:ident, $($args:expr),*) => {
-                if matches!($self.dtype(), DataType::Object(_)) {
+                if matches!($self.dtype(), DataType::Object(_, _)) {
                     // let ca = $self.0.unpack::<ObjectType<ObjectValue>>().unwrap();
                     // ca.$method($($args),*)
                     todo!()
@@ -463,7 +482,7 @@ impl RbSeries {
             DataType::Datetime(_, _)
                 | DataType::Date
                 | DataType::Duration(_)
-                | DataType::Categorical(_)
+                | DataType::Categorical(_, _)
                 | DataType::Time
         ) || !skip_nulls
         {
@@ -604,12 +623,12 @@ impl RbSeries {
                 )?;
                 ca.into_datetime(tu, tz).into_series()
             }
-            Some(DataType::Utf8) => {
+            Some(DataType::String) => {
                 let ca = dispatch_apply!(series, apply_lambda_with_utf8_out_type, lambda, 0, None)?;
 
                 ca.into_series()
             }
-            Some(DataType::Object(_)) => {
+            Some(DataType::Object(_, _)) => {
                 let ca =
                     dispatch_apply!(series, apply_lambda_with_object_out_type, lambda, 0, None)?;
                 ca.into_series()
@@ -656,8 +675,13 @@ impl RbSeries {
         self.series.borrow_mut().shrink_to_fit();
     }
 
-    pub fn dot(&self, other: &RbSeries) -> Option<f64> {
-        self.series.borrow().dot(&other.series.borrow())
+    pub fn dot(&self, other: &RbSeries) -> RbResult<f64> {
+        let out = self
+            .series
+            .borrow()
+            .dot(&other.series.borrow())
+            .map_err(RbPolarsErr::from)?;
+        Ok(out)
     }
 
     pub fn skew(&self, bias: bool) -> RbResult<Option<f64>> {

data/ext/polars/src/series/{set_at_idx.rs → scatter.rs}

@@ -5,9 +5,9 @@ use crate::error::RbPolarsErr;
 use crate::{RbResult, RbSeries};
 
 impl RbSeries {
-    pub fn set_at_idx(&self, idx: &RbSeries, values: &RbSeries) -> RbResult<()> {
+    pub fn scatter(&self, idx: &RbSeries, values: &RbSeries) -> RbResult<()> {
         let mut s = self.series.borrow_mut();
-        match set_at_idx(s.clone(), &idx.series.borrow(), &values.series.borrow()) {
+        match scatter(s.clone(), &idx.series.borrow(), &values.series.borrow()) {
             Ok(out) => {
                 *s = out;
                 Ok(())
@@ -17,7 +17,7 @@ impl RbSeries {
     }
 }
 
-fn set_at_idx(mut s: Series, idx: &Series, values: &Series) -> PolarsResult<Series> {
+fn scatter(mut s: Series, idx: &Series, values: &Series) -> PolarsResult<Series> {
     let logical_dtype = s.dtype().clone();
     let idx = idx.cast(&IDX_DTYPE)?;
     let idx = idx.rechunk();
@@ -43,62 +43,62 @@ fn set_at_idx(mut s: Series, idx: &Series, values: &Series) -> PolarsResult<Seri
         DataType::Int8 => {
             let ca: &mut ChunkedArray<Int8Type> = mutable_s.as_mut();
             let values = values.i8()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::Int16 => {
             let ca: &mut ChunkedArray<Int16Type> = mutable_s.as_mut();
             let values = values.i16()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::Int32 => {
             let ca: &mut ChunkedArray<Int32Type> = mutable_s.as_mut();
             let values = values.i32()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::Int64 => {
             let ca: &mut ChunkedArray<Int64Type> = mutable_s.as_mut();
             let values = values.i64()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::UInt8 => {
             let ca: &mut ChunkedArray<UInt8Type> = mutable_s.as_mut();
             let values = values.u8()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::UInt16 => {
             let ca: &mut ChunkedArray<UInt16Type> = mutable_s.as_mut();
             let values = values.u16()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::UInt32 => {
             let ca: &mut ChunkedArray<UInt32Type> = mutable_s.as_mut();
             let values = values.u32()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::UInt64 => {
             let ca: &mut ChunkedArray<UInt64Type> = mutable_s.as_mut();
             let values = values.u64()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::Float32 => {
             let ca: &mut ChunkedArray<Float32Type> = mutable_s.as_mut();
             let values = values.f32()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::Float64 => {
             let ca: &mut ChunkedArray<Float64Type> = mutable_s.as_mut();
             let values = values.f64()?;
-            std::mem::take(ca).set_at_idx2(idx, values)
+            std::mem::take(ca).scatter(idx, values)
         }
         DataType::Boolean => {
             let ca = s.bool()?;
             let values = values.bool()?;
-            ca.set_at_idx2(idx, values)
+            ca.scatter(idx, values)
         }
-        DataType::Utf8 => {
-            let ca = s.utf8()?;
-            let values = values.utf8()?;
-            ca.set_at_idx2(idx, values)
+        DataType::String => {
+            let ca = s.str()?;
+            let values = values.str()?;
+            ca.scatter(idx, values)
         }
         _ => panic!("not yet implemented for dtype: {}", logical_dtype),
     };

data/ext/polars/src/utils.rs

@@ -23,7 +23,7 @@ macro_rules! apply_method_all_arrow_series2 {
     ($self:expr, $method:ident, $($args:expr),*) => {
         match $self.dtype() {
             DataType::Boolean => $self.bool().unwrap().$method($($args),*),
-            DataType::Utf8 => $self.utf8().unwrap().$method($($args),*),
+            DataType::String => $self.str().unwrap().$method($($args),*),
             DataType::UInt8 => $self.u8().unwrap().$method($($args),*),
             DataType::UInt16 => $self.u16().unwrap().$method($($args),*),
             DataType::UInt32 => $self.u32().unwrap().$method($($args),*),