polars-df 0.23.0 → 0.24.0

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

@@ -77,21 +77,22 @@ pub(crate) fn get_rbseq(obj: Value) -> RbResult<(RArray, usize)> {
 
 pub(crate) fn get_df(obj: Value) -> RbResult<DataFrame> {
     let rbdf = obj.funcall::<_, _, &RbDataFrame>("_df", ())?;
-    Ok(rbdf.df.borrow().clone())
+    Ok(rbdf.df.read().clone())
 }
 
 pub(crate) fn get_lf(obj: Value) -> RbResult<LazyFrame> {
     let rbdf = obj.funcall::<_, _, &RbLazyFrame>("_ldf", ())?;
-    Ok(rbdf.ldf.borrow().clone())
+    Ok(rbdf.ldf.read().clone())
 }
 
 pub(crate) fn get_series(obj: Value) -> RbResult<Series> {
     let rbs = obj.funcall::<_, _, &RbSeries>("_s", ())?;
-    Ok(rbs.series.borrow().clone())
+    Ok(rbs.series.read().clone())
 }
 
 pub(crate) fn to_series(s: RbSeries) -> Value {
-    let series = pl_series();
+    let ruby = Ruby::get().unwrap();
+    let series = pl_series(&ruby);
     series
         .funcall::<_, _, Value>("_from_rbseries", (s,))
         .unwrap()
@@ -136,7 +137,7 @@ fn struct_dict<'a>(ruby: &Ruby, vals: impl Iterator<Item = AnyValue<'a>>, flds:
 
 impl IntoValue for Wrap<DataType> {
     fn into_value_with(self, ruby: &Ruby) -> Value {
-        let pl = crate::rb_modules::polars();
+        let pl = crate::rb_modules::polars(ruby);
 
         match self.0 {
             DataType::Int8 => {
@@ -338,10 +339,12 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::Int16" => DataType::Int16,
                 "Polars::Int32" => DataType::Int32,
                 "Polars::Int64" => DataType::Int64,
+                "Polars::Int128" => DataType::Int64,
                 "Polars::UInt8" => DataType::UInt8,
                 "Polars::UInt16" => DataType::UInt16,
                 "Polars::UInt32" => DataType::UInt32,
                 "Polars::UInt64" => DataType::UInt64,
+                "Polars::UInt128" => DataType::UInt128,
                 "Polars::Float32" => DataType::Float32,
                 "Polars::Float64" => DataType::Float64,
                 "Polars::Boolean" => DataType::Boolean,
@@ -367,12 +370,12 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::Object" => DataType::Object(OBJECT_NAME),
                 "Polars::Unknown" => DataType::Unknown(Default::default()),
                 dt => {
-                    return Err(RbValueError::new_err(format!(
-                        "{dt} is not a correct polars DataType.",
+                    return Err(RbTypeError::new_err(format!(
+                        "'{dt}' is not a Polars data type",
                     )));
                 }
             }
-        } else if String::try_convert(ob).is_err() {
+        } else {
             let cls = ob.class();
             let name = unsafe { cls.name() }.into_owned();
             match name.as_str() {
@@ -380,10 +383,12 @@ impl TryConvert for Wrap<DataType> {
                 "Polars::Int16" => DataType::Int16,
                 "Polars::Int32" => DataType::Int32,
                 "Polars::Int64" => DataType::Int64,
+                "Polars::Int128" => DataType::Int128,
                 "Polars::UInt8" => DataType::UInt8,
                 "Polars::UInt16" => DataType::UInt16,
                 "Polars::UInt32" => DataType::UInt32,
                 "Polars::UInt64" => DataType::UInt64,
+                "Polars::UInt128" => DataType::UInt128,
                 "Polars::Float32" => DataType::Float32,
                 "Polars::Float64" => DataType::Float64,
                 "Polars::Boolean" => DataType::Boolean,
@@ -448,42 +453,11 @@ impl TryConvert for Wrap<DataType> {
                     DataType::Struct(fields)
                 }
                 "Polars::Null" => DataType::Null,
-                "Object" => DataType::Object(OBJECT_NAME),
+                "Polars::Object" => DataType::Object(OBJECT_NAME),
                 "Polars::Unknown" => DataType::Unknown(Default::default()),
                 dt => {
                     return Err(RbTypeError::new_err(format!(
-                        "A {dt} object is not a correct polars DataType. \
-                        Hint: use the class without instantiating it.",
-                    )));
-                }
-            }
-        } else {
-            match String::try_convert(ob)?.as_str() {
-                "u8" => DataType::UInt8,
-                "u16" => DataType::UInt16,
-                "u32" => DataType::UInt32,
-                "u64" => DataType::UInt64,
-                "i8" => DataType::Int8,
-                "i16" => DataType::Int16,
-                "i32" => DataType::Int32,
-                "i64" => DataType::Int64,
-                "str" => DataType::String,
-                "bin" => DataType::Binary,
-                "bool" => DataType::Boolean,
-                "cat" => DataType::from_categories(Categories::global()),
-                "date" => DataType::Date,
-                "datetime" => DataType::Datetime(TimeUnit::Microseconds, None),
-                "f32" => DataType::Float32,
-                "time" => DataType::Time,
-                "dur" => DataType::Duration(TimeUnit::Microseconds),
-                "f64" => DataType::Float64,
-                "obj" => DataType::Object(OBJECT_NAME),
-                "list" => DataType::List(Box::new(DataType::Boolean)),
-                "null" => DataType::Null,
-                "unk" => DataType::Unknown(Default::default()),
-                _ => {
-                    return Err(RbValueError::new_err(format!(
-                        "{ob} is not a supported DataType."
+                        "'{dt}' is not a Polars data type",
                     )));
                 }
             }
@@ -531,7 +505,7 @@ impl<'s> TryConvert for Wrap<Row<'s>> {
 
 impl TryConvert for Wrap<Schema> {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        let dict = RHash::try_convert(ob)?;
+        let dict: RHash = ob.funcall("to_h", ())?;
 
         let mut schema = Vec::new();
         dict.foreach(|key: String, val: Wrap<DataType>| {
@@ -628,7 +602,7 @@ impl TryConvert for Wrap<ScanSources> {
             }
             RubyScanSourceInput::File(file) => {
                 let mut sources = Vec::with_capacity(num_items);
-                sources.push(file);
+                sources.push(file.into());
                 MutableSources::Files(sources)
             }
             RubyScanSourceInput::Buffer(buffer) => {
@@ -641,7 +615,7 @@ impl TryConvert for Wrap<ScanSources> {
         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::Files(v), RubyScanSourceInput::File(f)) => v.push(f.into()),
                 (MutableSources::Buffers(v), RubyScanSourceInput::Buffer(f)) => v.push(f),
                 _ => {
                     return Err(RbTypeError::new_err(
@@ -1052,6 +1026,24 @@ impl TryConvert for Wrap<RankMethod> {
     }
 }
 
+impl TryConvert for Wrap<RollingRankMethod> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match String::try_convert(ob)?.as_str() {
+            "min" => RollingRankMethod::Min,
+            "max" => RollingRankMethod::Max,
+            "average" => RollingRankMethod::Average,
+            "dense" => RollingRankMethod::Dense,
+            "random" => RollingRankMethod::Random,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "rank `method` must be one of {{'min', 'max', 'average', 'dense', 'random'}}, got {v}",
+                )));
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<Roll> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match String::try_convert(ob)?.as_str() {
@@ -1091,9 +1083,11 @@ impl TryConvert for Wrap<UniqueKeepStrategy> {
         let parsed = match String::try_convert(ob)?.as_str() {
             "first" => UniqueKeepStrategy::First,
             "last" => UniqueKeepStrategy::Last,
+            "none" => UniqueKeepStrategy::None,
+            "any" => UniqueKeepStrategy::Any,
             v => {
                 return Err(RbValueError::new_err(format!(
-                    "keep must be one of {{'first', 'last'}}, got {v}"
+                    "`keep` must be one of {{'first', 'last', 'any', 'none'}}, got {v}",
                 )));
             }
         };
@@ -1467,6 +1461,17 @@ impl TryConvert for Wrap<UnicodeForm> {
     }
 }
 
+impl TryConvert for Wrap<Option<KeyValueMetadata>> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        if ob.is_nil() {
+            return Ok(Wrap(None));
+        }
+        todo!();
+    }
+}
+
+unsafe impl TryConvertOwned for Wrap<Option<KeyValueMetadata>> {}
+
 impl TryConvert for Wrap<Option<TimeZone>> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let tz = Option::<Wrap<PlSmallStr>>::try_convert(ob)?;

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

@@ -3,19 +3,19 @@ use magnus::{IntoValue, Ruby, Value, prelude::*};
 use super::*;
 use crate::RbResult;
 use crate::conversion::{ObjectValue, Wrap};
-use crate::interop::arrow::to_ruby::dataframe_to_stream;
+use crate::interop::arrow::to_rb::dataframe_to_stream;
 
 impl RbDataFrame {
-    pub fn row_tuple(ruby: &Ruby, rb_self: &Self, idx: i64) -> Value {
+    pub fn row_tuple(ruby: &Ruby, self_: &Self, idx: i64) -> Value {
         let idx = if idx < 0 {
-            (rb_self.df.borrow().height() as i64 + idx) as usize
+            (self_.df.read().height() as i64 + idx) as usize
         } else {
             idx as usize
         };
         ruby.ary_from_iter(
-            rb_self
+            self_
                 .df
-                .borrow()
+                .read()
                 .get_columns()
                 .iter()
                 .map(|s| match s.dtype() {
@@ -29,13 +29,13 @@ impl RbDataFrame {
         .as_value()
     }
 
-    pub fn row_tuples(ruby: &Ruby, rb_self: &Self) -> Value {
-        let df = &rb_self.df;
-        ruby.ary_from_iter((0..df.borrow().height()).map(|idx| {
+    pub fn row_tuples(ruby: &Ruby, self_: &Self) -> Value {
+        let df = &self_.df;
+        ruby.ary_from_iter((0..df.read().height()).map(|idx| {
             ruby.ary_from_iter(
-                rb_self
+                self_
                     .df
-                    .borrow()
+                    .read()
                     .get_columns()
                     .iter()
                     .map(|s| match s.dtype() {
@@ -50,8 +50,8 @@ impl RbDataFrame {
         .as_value()
     }
 
-    pub fn __arrow_c_stream__(ruby: &Ruby, rb_self: &Self) -> RbResult<Value> {
-        rb_self.df.borrow_mut().align_chunks();
-        dataframe_to_stream(&rb_self.df.borrow(), ruby)
+    pub fn __arrow_c_stream__(ruby: &Ruby, self_: &Self) -> RbResult<Value> {
+        self_.df.write().align_chunks();
+        dataframe_to_stream(&self_.df.read(), ruby)
     }
 }