polars-df 0.1.3 → 0.1.5

data/ext/polars/src/conversion.rs

@@ -1,33 +1,58 @@
-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())
 }
@@ -87,6 +112,25 @@ impl From<Wrap<AnyValue<'_>>> for Value {
                 .unwrap()
                 .funcall::<_, _, Value>("to_date", ())
                 .unwrap(),
+            AnyValue::Datetime(v, tu, tz) => {
+                let t = match tu {
+                    TimeUnit::Nanoseconds => todo!(),
+                    TimeUnit::Microseconds => {
+                        let sec = v / 1000000;
+                        let subsec = v % 1000000;
+                        class::time()
+                            .funcall::<_, _, Value>("at", (sec, subsec, Symbol::new("usec")))
+                            .unwrap()
+                    }
+                    TimeUnit::Milliseconds => todo!(),
+                };
+
+                if tz.is_some() {
+                    todo!();
+                } else {
+                    t.funcall::<_, _, Value>("utc", ()).unwrap()
+                }
+            }
             _ => todo!(),
         }
     }
@@ -150,6 +194,39 @@ impl<'s> TryConvert for Wrap<AnyValue<'s>> {
     }
 }
 
+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() {
@@ -238,6 +315,22 @@ impl TryConvert for Wrap<JoinType> {
     }
 }
 
+impl TryConvert for Wrap<ListToStructWidthStrategy> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let parsed = match ob.try_convert::<String>()?.as_str() {
+            "first_non_null" => ListToStructWidthStrategy::FirstNonNull,
+            "max_width" => ListToStructWidthStrategy::MaxWidth,
+            v => {
+                return Err(RbValueError::new_err(format!(
+                    "n_field_strategy must be one of {{'first_non_null', 'max_width'}}, got {}",
+                    v
+                )))
+            }
+        };
+        Ok(Wrap(parsed))
+    }
+}
+
 impl TryConvert for Wrap<NullBehavior> {
     fn try_convert(ob: Value) -> RbResult<Self> {
         let parsed = match ob.try_convert::<String>()?.as_str() {
@@ -425,18 +518,100 @@ pub fn parse_parquet_compression(
     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,
 }
 
+impl Hash for ObjectValue {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        let h = self
+            .inner
+            .funcall::<_, _, isize>("hash", ())
+            .expect("should be hashable");
+        state.write_isize(h)
+    }
+}
+
+impl Eq for ObjectValue {}
+
+impl PartialEq for ObjectValue {
+    fn eq(&self, other: &Self) -> bool {
+        self.inner.eql(&other.inner).unwrap_or(false)
+    }
+}
+
+impl Display for ObjectValue {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.inner)
+    }
+}
+
+impl PolarsObject for ObjectValue {
+    fn type_name() -> &'static str {
+        "object"
+    }
+}
+
+impl From<Value> for ObjectValue {
+    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 {
+        ObjectValue { inner: *QNIL }
+    }
+}

data/ext/polars/src/dataframe.rs

@@ -1,15 +1,21 @@
 use magnus::{r_hash::ForEach, RArray, RHash, RString, Value};
+use polars::io::avro::AvroCompression;
 use polars::io::mmap::ReaderBytes;
 use polars::io::RowCount;
+use polars::prelude::pivot::{pivot, pivot_stable};
 use polars::prelude::*;
 use std::cell::RefCell;
 use std::io::{BufWriter, Cursor};
 use std::ops::Deref;
 
+use crate::apply::dataframe::{
+    apply_lambda_unknown, apply_lambda_with_bool_out_type, apply_lambda_with_primitive_out_type,
+    apply_lambda_with_utf8_out_type,
+};
 use crate::conversion::*;
 use crate::file::{get_file_like, get_mmap_bytes_reader};
 use crate::series::{to_rbseries_collection, to_series_collection};
-use crate::{series, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
+use crate::{series, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries};
 
 #[magnus::wrap(class = "Polars::RbDataFrame")]
 pub struct RbDataFrame {
@@ -179,6 +185,48 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
+    pub fn read_avro(
+        rb_f: Value,
+        columns: Option<Vec<String>>,
+        projection: Option<Vec<usize>>,
+        n_rows: Option<usize>,
+    ) -> RbResult<Self> {
+        use polars::io::avro::AvroReader;
+
+        let file = get_file_like(rb_f, false)?;
+        let df = AvroReader::new(file)
+            .with_projection(projection)
+            .with_columns(columns)
+            .with_n_rows(n_rows)
+            .finish()
+            .map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
+    pub fn write_avro(
+        &self,
+        rb_f: Value,
+        compression: Wrap<Option<AvroCompression>>,
+    ) -> RbResult<()> {
+        use polars::io::avro::AvroWriter;
+
+        if let Ok(s) = rb_f.try_convert::<String>() {
+            let f = std::fs::File::create(&s).unwrap();
+            AvroWriter::new(f)
+                .with_compression(compression.0)
+                .finish(&mut self.df.borrow_mut())
+                .map_err(RbPolarsErr::from)?;
+        } else {
+            let mut buf = get_file_like(rb_f, true)?;
+            AvroWriter::new(&mut buf)
+                .with_compression(compression.0)
+                .finish(&mut self.df.borrow_mut())
+                .map_err(RbPolarsErr::from)?;
+        }
+
+        Ok(())
+    }
+
     pub fn read_json(rb_f: Value) -> RbResult<Self> {
         // memmap the file first
         let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
@@ -238,6 +286,14 @@ impl RbDataFrame {
         Ok(())
     }
 
+    pub fn read_hashes(
+        _dicts: Value,
+        _infer_schema_length: Option<usize>,
+        _schema_overwrite: Option<Wrap<Schema>>,
+    ) -> RbResult<Self> {
+        Err(RbPolarsErr::todo())
+    }
+
     pub fn read_hash(data: RHash) -> RbResult<Self> {
         let mut cols: Vec<Series> = Vec::new();
         data.foreach(|name: String, values: Value| {
@@ -751,6 +807,31 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
+    pub fn pivot_expr(
+        &self,
+        values: Vec<String>,
+        index: Vec<String>,
+        columns: Vec<String>,
+        aggregate_expr: &RbExpr,
+        maintain_order: bool,
+        sort_columns: bool,
+    ) -> RbResult<Self> {
+        let fun = match maintain_order {
+            true => pivot_stable,
+            false => pivot,
+        };
+        let df = fun(
+            &self.df.borrow(),
+            values,
+            index,
+            columns,
+            aggregate_expr.inner.clone(),
+            sort_columns,
+        )
+        .map_err(RbPolarsErr::from)?;
+        Ok(RbDataFrame::new(df))
+    }
+
     pub fn partition_by(&self, groups: Vec<String>, stable: bool) -> RbResult<Vec<Self>> {
         let out = if stable {
             self.df.borrow().partition_by_stable(groups)
@@ -870,10 +951,74 @@ impl RbDataFrame {
         df.into()
     }
 
+    pub fn apply(
+        &self,
+        lambda: Value,
+        output_type: Option<Wrap<DataType>>,
+        inference_size: usize,
+    ) -> RbResult<(Value, bool)> {
+        let df = &self.df.borrow();
+
+        let output_type = output_type.map(|dt| dt.0);
+        let out = match output_type {
+            Some(DataType::Int32) => {
+                apply_lambda_with_primitive_out_type::<Int32Type>(df, lambda, 0, None).into_series()
+            }
+            Some(DataType::Int64) => {
+                apply_lambda_with_primitive_out_type::<Int64Type>(df, lambda, 0, None).into_series()
+            }
+            Some(DataType::UInt32) => {
+                apply_lambda_with_primitive_out_type::<UInt32Type>(df, lambda, 0, None)
+                    .into_series()
+            }
+            Some(DataType::UInt64) => {
+                apply_lambda_with_primitive_out_type::<UInt64Type>(df, lambda, 0, None)
+                    .into_series()
+            }
+            Some(DataType::Float32) => {
+                apply_lambda_with_primitive_out_type::<Float32Type>(df, lambda, 0, None)
+                    .into_series()
+            }
+            Some(DataType::Float64) => {
+                apply_lambda_with_primitive_out_type::<Float64Type>(df, lambda, 0, None)
+                    .into_series()
+            }
+            Some(DataType::Boolean) => {
+                apply_lambda_with_bool_out_type(df, lambda, 0, None).into_series()
+            }
+            Some(DataType::Date) => {
+                apply_lambda_with_primitive_out_type::<Int32Type>(df, lambda, 0, None)
+                    .into_date()
+                    .into_series()
+            }
+            Some(DataType::Datetime(tu, tz)) => {
+                apply_lambda_with_primitive_out_type::<Int64Type>(df, lambda, 0, None)
+                    .into_datetime(tu, tz)
+                    .into_series()
+            }
+            Some(DataType::Utf8) => {
+                apply_lambda_with_utf8_out_type(df, lambda, 0, None).into_series()
+            }
+            _ => return apply_lambda_unknown(df, lambda, inference_size),
+        };
+
+        Ok((RbSeries::from(out).into(), false))
+    }
+
     pub fn shrink_to_fit(&self) {
         self.df.borrow_mut().shrink_to_fit();
     }
 
+    pub fn hash_rows(&self, k0: u64, k1: u64, k2: u64, k3: u64) -> RbResult<RbSeries> {
+        let hb = ahash::RandomState::with_seeds(k0, k1, k2, k3);
+        let hash = self
+            .df
+            .borrow_mut()
+            .hash_rows(Some(hb))
+            .map_err(RbPolarsErr::from)?;
+        Ok(hash.into_series().into())
+    }
+
     pub fn transpose(&self, include_header: bool, names: String) -> RbResult<Self> {
         let mut df = self.df.borrow().transpose().map_err(RbPolarsErr::from)?;
         if include_header {

data/ext/polars/src/error.rs

@@ -22,6 +22,10 @@ impl RbPolarsErr {
     pub fn other(message: String) -> Error {
         Error::runtime_error(message)
     }
+
+    pub fn todo() -> Error {
+        Error::runtime_error("not implemented yet")
+    }
 }
 
 pub struct RbValueError {}
@@ -31,3 +35,11 @@ impl RbValueError {
         Error::new(arg_error(), message)
     }
 }
+
+pub struct ComputeError {}
+
+impl ComputeError {
+    pub fn new_err(message: String) -> Error {
+        Error::runtime_error(message)
+    }
+}

data/ext/polars/src/lazy/apply.rs

@@ -1,7 +1,39 @@
 use magnus::Value;
-use polars::error::PolarsResult;
-use polars::series::Series;
+use polars::prelude::*;
+
+use crate::lazy::dsl::RbExpr;
+use crate::Wrap;
 
 pub fn binary_lambda(_lambda: Value, _a: Series, _b: Series) -> PolarsResult<Series> {
     todo!();
 }
+
+pub fn map_single(
+    rbexpr: &RbExpr,
+    _lambda: Value,
+    output_type: Option<Wrap<DataType>>,
+    agg_list: bool,
+) -> RbExpr {
+    let output_type = output_type.map(|wrap| wrap.0);
+
+    let output_type2 = output_type.clone();
+    let function = move |_s: Series| {
+        let _output_type = output_type2.clone().unwrap_or(DataType::Unknown);
+
+        todo!();
+    };
+
+    let output_map = GetOutput::map_field(move |fld| match output_type {
+        Some(ref dt) => Field::new(fld.name(), dt.clone()),
+        None => {
+            let mut fld = fld.clone();
+            fld.coerce(DataType::Unknown);
+            fld
+        }
+    });
+    if agg_list {
+        rbexpr.clone().inner.map_list(function, output_map).into()
+    } else {
+        rbexpr.clone().inner.map(function, output_map).into()
+    }
+}

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

@@ -3,7 +3,7 @@ use polars::io::RowCount;
 use polars::lazy::frame::{LazyFrame, LazyGroupBy};
 use polars::prelude::*;
 use std::cell::RefCell;
-use std::io::BufWriter;
+use std::io::{BufWriter, Read};
 
 use crate::conversion::*;
 use crate::file::get_file_like;
@@ -53,6 +53,27 @@ impl From<LazyFrame> for RbLazyFrame {
 }
 
 impl RbLazyFrame {
+    pub fn read_json(rb_f: Value) -> RbResult<Self> {
+        // it is faster to first read to memory and then parse: https://github.com/serde-rs/json/issues/160
+        // so don't bother with files.
+        let mut json = String::new();
+        let _ = get_file_like(rb_f, false)?
+            .read_to_string(&mut json)
+            .unwrap();
+
+        // Safety
+        // we skipped the serializing/deserializing of the static in lifetime in `DataType`
+        // so we actually don't have a lifetime at all when serializing.
+
+        // &str still has a lifetime. Bit its ok, because we drop it immediately
+        // in this scope
+        let json = unsafe { std::mem::transmute::<&'_ str, &'static str>(json.as_str()) };
+
+        let lp = serde_json::from_str::<LogicalPlan>(json)
+            .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
+        Ok(LazyFrame::from(lp).into())
+    }
+
     pub fn new_from_ndjson(
         path: String,
         infer_schema_length: Option<usize>,
@@ -211,7 +232,7 @@ impl RbLazyFrame {
         projection_pushdown: bool,
         simplify_expr: bool,
         slice_pushdown: bool,
-        _cse: bool,
+        cse: bool,
         allow_streaming: bool,
     ) -> RbLazyFrame {
         let ldf = self.ldf.clone();
@@ -220,7 +241,7 @@ impl RbLazyFrame {
             .with_predicate_pushdown(predicate_pushdown)
             .with_simplify_expr(simplify_expr)
             .with_slice_pushdown(slice_pushdown)
-            // .with_common_subplan_elimination(cse)
+            .with_common_subplan_elimination(cse)
             .with_streaming(allow_streaming)
             .with_projection_pushdown(projection_pushdown);
         ldf.into()
@@ -349,6 +370,56 @@ impl RbLazyFrame {
         })
     }
 
+    pub fn with_context(&self, contexts: RArray) -> RbResult<Self> {
+        let contexts = contexts
+            .each()
+            .map(|v| v.unwrap().try_convert())
+            .collect::<RbResult<Vec<&RbLazyFrame>>>()?;
+        let contexts = contexts
+            .into_iter()
+            .map(|ldf| ldf.ldf.clone())
+            .collect::<Vec<_>>();
+        Ok(self.ldf.clone().with_context(contexts).into())
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn join_asof(
+        &self,
+        other: &RbLazyFrame,
+        left_on: &RbExpr,
+        right_on: &RbExpr,
+        left_by: Option<Vec<String>>,
+        right_by: Option<Vec<String>>,
+        allow_parallel: bool,
+        force_parallel: bool,
+        suffix: String,
+        strategy: Wrap<AsofStrategy>,
+        tolerance: Option<Wrap<AnyValue<'_>>>,
+        tolerance_str: Option<String>,
+    ) -> RbResult<Self> {
+        let ldf = self.ldf.clone();
+        let other = other.ldf.clone();
+        let left_on = left_on.inner.clone();
+        let right_on = right_on.inner.clone();
+        Ok(ldf
+            .join_builder()
+            .with(other)
+            .left_on([left_on])
+            .right_on([right_on])
+            .allow_parallel(allow_parallel)
+            .force_parallel(force_parallel)
+            .how(JoinType::AsOf(AsOfOptions {
+                strategy: strategy.0,
+                left_by,
+                right_by,
+                tolerance: tolerance.map(|t| t.0.into_static().unwrap()),
+                tolerance_str,
+            }))
+            .suffix(suffix)
+            .finish()
+            .into())
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn join(
         &self,