polars-df 0.15.0 → 0.26.0

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

@@ -1,5 +1,5 @@
-use magnus::{prelude::*, r_hash::ForEach, RArray, RHash, Symbol, Value};
-use polars::frame::row::{rows_to_schema_supertypes, rows_to_supertypes, Row};
+use magnus::{RArray, RHash, Ruby, Symbol, Value, prelude::*, r_hash::ForEach};
+use polars::frame::row::{Row, rows_to_schema_supertypes, rows_to_supertypes};
 use polars::prelude::*;
 
 use super::*;
@@ -54,7 +54,7 @@ fn finish_from_rows(
     schema_overrides: Option<Schema>,
     infer_schema_length: Option<usize>,
 ) -> RbResult<RbDataFrame> {
-    let mut schema = if let Some(mut schema) = schema {
+    let schema = if let Some(mut schema) = schema {
         resolve_schema_overrides(&mut schema, schema_overrides);
         update_schema_from_rows(&mut schema, &rows, infer_schema_length)?;
         schema
@@ -62,11 +62,6 @@ fn finish_from_rows(
         rows_to_schema_supertypes(&rows, infer_schema_length).map_err(RbPolarsErr::from)?
     };
 
-    // TODO: Remove this step when Decimals are supported properly.
-    // Erasing the decimal precision/scale here will just require us to infer it again later.
-    // https://github.com/pola-rs/polars/issues/14427
-    erase_decimal_precision_scale(&mut schema);
-
     let df = DataFrame::from_rows_and_schema(&rows, &schema).map_err(RbPolarsErr::from)?;
     Ok(df.into())
 }
@@ -106,14 +101,6 @@ fn resolve_schema_overrides(schema: &mut Schema, schema_overrides: Option<Schema
     }
 }
 
-fn erase_decimal_precision_scale(schema: &mut Schema) {
-    for dtype in schema.iter_values_mut() {
-        if let DataType::Decimal(_, _) = dtype {
-            *dtype = DataType::Decimal(None, None)
-        }
-    }
-}
-
 fn columns_names_to_empty_schema<'a, I>(column_names: I) -> Schema
 where
     I: IntoIterator<Item = &'a str>,
@@ -125,6 +112,7 @@ where
 }
 
 fn dicts_to_rows<'a>(data: &Value, names: &'a [String], _strict: bool) -> RbResult<Vec<Row<'a>>> {
+    let ruby = Ruby::get_with(*data);
     let (data, len) = get_rbseq(*data)?;
     let mut rows = Vec::with_capacity(len);
     for d in data.into_iter() {
@@ -132,8 +120,8 @@ fn dicts_to_rows<'a>(data: &Value, names: &'a [String], _strict: bool) -> RbResu
 
         let mut row = Vec::with_capacity(names.len());
         for k in names.iter() {
-            // TODO improve performance
-            let val = match d.get(k.clone()).or_else(|| d.get(Symbol::new(k))) {
+            // TODO improve performance (must work with GC)
+            let val = match d.get(k.clone()).or_else(|| d.get(ruby.to_symbol(k))) {
                 None => AnyValue::Null,
                 Some(val) => Wrap::<AnyValue>::try_convert(val)?.0,
             };

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

@@ -1,55 +1,46 @@
-use magnus::{prelude::*, IntoValue, RArray, Value};
+use magnus::{Ruby, Value, prelude::*};
 
 use super::*;
-use crate::conversion::{ObjectValue, Wrap};
-use crate::interop::arrow::to_ruby::dataframe_to_stream;
 use crate::RbResult;
+use crate::conversion::{ObjectValue, Wrap};
+use crate::interop::arrow::to_rb::dataframe_to_stream;
+use crate::ruby::utils::TryIntoValue;
 
 impl RbDataFrame {
-    pub fn row_tuple(&self, idx: i64) -> Value {
+    pub fn row_tuple(ruby: &Ruby, self_: &Self, idx: i64) -> RbResult<Value> {
         let idx = if idx < 0 {
-            (self.df.borrow().height() as i64 + idx) as usize
+            (self_.df.read().height() as i64 + idx) as usize
         } else {
             idx as usize
         };
-        RArray::from_iter(
-            self.df
-                .borrow()
-                .get_columns()
-                .iter()
-                .map(|s| match s.dtype() {
-                    DataType::Object(_, _) => {
-                        let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
-                        obj.unwrap().to_value()
-                    }
-                    _ => Wrap(s.get(idx).unwrap()).into_value(),
-                }),
-        )
-        .as_value()
+        ruby.ary_try_from_iter(self_.df.read().columns().iter().map(|s| match s.dtype() {
+            DataType::Object(_) => {
+                let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
+                // TODO remove unwrap and clone
+                obj.unwrap().clone().try_into_value_with(ruby)
+            }
+            _ => Wrap(s.get(idx).unwrap()).try_into_value_with(ruby),
+        }))
+        .map(|v| v.as_value())
     }
 
-    pub fn row_tuples(&self) -> Value {
-        let df = &self.df;
-        RArray::from_iter((0..df.borrow().height()).map(|idx| {
-            RArray::from_iter(
-                self.df
-                    .borrow()
-                    .get_columns()
-                    .iter()
-                    .map(|s| match s.dtype() {
-                        DataType::Object(_, _) => {
-                            let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
-                            obj.unwrap().to_value()
-                        }
-                        _ => Wrap(s.get(idx).unwrap()).into_value(),
-                    }),
-            )
+    pub fn row_tuples(ruby: &Ruby, self_: &Self) -> RbResult<Value> {
+        let df = &self_.df;
+        ruby.ary_try_from_iter((0..df.read().height()).map(|idx| {
+            ruby.ary_try_from_iter(self_.df.read().columns().iter().map(|s| match s.dtype() {
+                DataType::Object(_) => {
+                    let obj: Option<&ObjectValue> = s.get_object(idx).map(|any| any.into());
+                    // TODO remove unwrap and clone
+                    obj.unwrap().clone().try_into_value_with(ruby)
+                }
+                _ => Wrap(s.get(idx).unwrap()).try_into_value_with(ruby),
+            }))
         }))
-        .as_value()
+        .map(|v| v.as_value())
     }
 
-    pub fn __arrow_c_stream__(&self) -> RbResult<Value> {
-        self.df.borrow_mut().align_chunks();
-        dataframe_to_stream(&self.df.borrow())
+    pub fn __arrow_c_stream__(ruby: &Ruby, self_: &Self) -> RbResult<Value> {
+        self_.df.write().align_chunks();
+        dataframe_to_stream(&self_.df.read(), ruby)
     }
 }