ext/polars/src/conversion.rs in polars-df-0.1.4 vs ext/polars/src/conversion.rs in polars-df-0.1.5

@@ -1,37 +1,60 @@
-use magnus::{class, RArray, Symbol, TryConvert, Value, QNIL};
+use magnus::{class, r_hash::ForEach, RArray, RHash, Symbol, TryConvert, Value, QNIL};
 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::DataFrame;
+use polars::io::avro::AvroCompression;
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 use std::fmt::{Display, Formatter};
 use std::hash::{Hash, Hasher};
 
-use crate::{RbDataFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
+use crate::{RbDataFrame, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
+pub(crate) fn slice_to_wrapped<T>(slice: &[T]) -> &[Wrap<T>] {
+    // Safety:
+    // Wrap is transparent.
+    unsafe { std::mem::transmute(slice) }
+}
+
+#[repr(transparent)]
 pub struct Wrap<T>(pub T);
 
+impl<T> Clone for Wrap<T>
+where
+    T: Clone,
+{
+    fn clone(&self) -> Self {
+        Wrap(self.0.clone())
+    }
+}
+
 impl<T> From<T> for Wrap<T> {
     fn from(t: T) -> Self {
         Wrap(t)
     }
 }
 
-pub fn get_rbseq(obj: Value) -> RbResult<(RArray, usize)> {
+pub(crate) fn get_rbseq(obj: Value) -> RbResult<(RArray, usize)> {
     let seq: RArray = obj.try_convert()?;
     let len = seq.len();
     Ok((seq, len))
 }
 
-pub fn get_df(obj: Value) -> RbResult<DataFrame> {
+pub(crate) fn get_df(obj: Value) -> RbResult<DataFrame> {
     let rbdf = obj.funcall::<_, _, &RbDataFrame>("_df", ())?;
     Ok(rbdf.df.borrow().clone())
 }
 
-pub fn get_series(obj: Value) -> RbResult<Series> {
+pub(crate) fn get_lf(obj: Value) -> RbResult<LazyFrame> {
+    let rbdf = obj.funcall::<_, _, &RbLazyFrame>("_ldf", ())?;
+    Ok(rbdf.ldf.clone())
+}
+
+pub(crate) fn get_series(obj: Value) -> RbResult<Series> {
     let rbs = obj.funcall::<_, _, &RbSeries>("_s", ())?;
     Ok(rbs.series.borrow().clone())
 }
 
 impl TryConvert for Wrap<Utf8Chunked> {
@@ -169,10 +192,43 @@
             )))
         }
     }
 }
 
+impl TryConvert for Wrap<AsofStrategy> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "backward" => AsofStrategy::Backward,
+            "forward" => AsofStrategy::Forward,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "strategy must be one of {{'backward', 'forward'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
+impl TryConvert for Wrap<Option<AvroCompression>> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "uncompressed" => None,
+            "snappy" => Some(AvroCompression::Snappy),
+            "deflate" => Some(AvroCompression::Deflate),
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "compression must be one of {{'uncompressed', 'snappy', 'deflate'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<CategoricalOrdering> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match ob.try_convert::<String>()?.as_str() {
             "physical" => CategoricalOrdering::Physical,
             "lexical" => CategoricalOrdering::Lexical,
@@ -460,10 +516,36 @@
         }
     };
     Ok(parsed)
 }
 
+impl<'s> TryConvert for Wrap<Row<'s>> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let mut vals: Vec<Wrap<AnyValue<'s>>> = Vec::new();
+        for item in ob.try_convert::<RArray>()?.each() {
+            vals.push(item?.try_convert::<Wrap<AnyValue<'s>>>()?);
+        }
+        let vals: Vec<AnyValue> = unsafe { std::mem::transmute(vals) };
+        Ok(Wrap(Row(vals)))
+    }
+}
+
+impl TryConvert for Wrap<Schema> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let dict = ob.try_convert::<RHash>()?;
+
+        let mut schema = Vec::new();
+        dict.foreach(|key: String, val: Wrap<DataType>| {
+            schema.push(Field::new(&key, val.0));
+            Ok(ForEach::Continue)
+        })
+        .unwrap();
+
+        Ok(Wrap(schema.into_iter().into()))
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct ObjectValue {
     pub inner: Value,
 }
 
@@ -501,18 +583,31 @@
     fn from(v: Value) -> Self {
         Self { inner: v }
     }
 }
 
+impl TryConvert for ObjectValue {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        Ok(ObjectValue { inner: ob })
+    }
+}
+
 impl From<&dyn PolarsObjectSafe> for &ObjectValue {
     fn from(val: &dyn PolarsObjectSafe) -> Self {
         unsafe { &*(val as *const dyn PolarsObjectSafe as *const ObjectValue) }
     }
 }
 
+// TODO remove
 impl ObjectValue {
     pub fn to_object(&self) -> Value {
         self.inner
+    }
+}
+
+impl From<ObjectValue> for Value {
+    fn from(val: ObjectValue) -> Self {
+        val.inner
     }
 }
 
 impl Default for ObjectValue {
     fn default() -> Self {