polars-df 0.4.0 → 0.6.0

data/ext/polars/src/expr/list.rs

@@ -0,0 +1,151 @@
+use magnus::Value;
+use polars::lazy::dsl::lit;
+use polars::prelude::*;
+use polars::series::ops::NullBehavior;
+
+use crate::conversion::Wrap;
+use crate::{RbExpr, RbResult};
+
+impl RbExpr {
+    pub fn list_arg_max(&self) -> Self {
+        self.inner.clone().list().arg_max().into()
+    }
+
+    pub fn list_arg_min(&self) -> Self {
+        self.inner.clone().list().arg_min().into()
+    }
+
+    pub fn list_contains(&self, other: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .contains(other.inner.clone())
+            .into()
+    }
+
+    pub fn list_count_match(&self, expr: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .count_match(expr.inner.clone())
+            .into()
+    }
+
+    pub fn list_diff(&self, n: i64, null_behavior: Wrap<NullBehavior>) -> RbResult<Self> {
+        Ok(self.inner.clone().list().diff(n, null_behavior.0).into())
+    }
+
+    pub fn list_eval(&self, expr: &RbExpr, parallel: bool) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .eval(expr.inner.clone(), parallel)
+            .into()
+    }
+
+    pub fn list_get(&self, index: &RbExpr) -> Self {
+        self.inner.clone().list().get(index.inner.clone()).into()
+    }
+
+    pub fn list_join(&self, separator: String) -> Self {
+        self.inner.clone().list().join(&separator).into()
+    }
+
+    pub fn list_lengths(&self) -> Self {
+        self.inner.clone().list().lengths().into()
+    }
+
+    pub fn list_max(&self) -> Self {
+        self.inner.clone().list().max().into()
+    }
+
+    pub fn list_mean(&self) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .mean()
+            .with_fmt("list.mean")
+            .into()
+    }
+
+    pub fn list_min(&self) -> Self {
+        self.inner.clone().list().min().into()
+    }
+
+    pub fn list_reverse(&self) -> Self {
+        self.inner.clone().list().reverse().into()
+    }
+
+    pub fn list_shift(&self, periods: i64) -> Self {
+        self.inner.clone().list().shift(periods).into()
+    }
+
+    pub fn list_slice(&self, offset: &RbExpr, length: Option<&RbExpr>) -> Self {
+        let length = match length {
+            Some(i) => i.inner.clone(),
+            None => lit(i64::MAX),
+        };
+        self.inner
+            .clone()
+            .list()
+            .slice(offset.inner.clone(), length)
+            .into()
+    }
+
+    pub fn list_sort(&self, reverse: bool) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .sort(SortOptions {
+                descending: reverse,
+                ..Default::default()
+            })
+            .with_fmt("list.sort")
+            .into()
+    }
+
+    pub fn list_sum(&self) -> Self {
+        self.inner.clone().list().sum().with_fmt("list.sum").into()
+    }
+
+    pub fn list_take(&self, index: &RbExpr, null_on_oob: bool) -> Self {
+        self.inner
+            .clone()
+            .list()
+            .take(index.inner.clone(), null_on_oob)
+            .into()
+    }
+
+    pub fn list_to_struct(
+        &self,
+        width_strat: Wrap<ListToStructWidthStrategy>,
+        _name_gen: Option<Value>,
+        upper_bound: usize,
+    ) -> RbResult<Self> {
+        // TODO fix
+        let name_gen = None;
+        // let name_gen = name_gen.map(|lambda| {
+        //     Arc::new(move |idx: usize| {
+        //         let out: Value = lambda.funcall("call", (idx,)).unwrap();
+        //         out.try_convert::<String>().unwrap()
+        //     }) as NameGenerator
+        // });
+
+        Ok(self
+            .inner
+            .clone()
+            .list()
+            .to_struct(width_strat.0, name_gen, upper_bound)
+            .into())
+    }
+
+    pub fn list_unique(&self, maintain_order: bool) -> Self {
+        let e = self.inner.clone();
+
+        if maintain_order {
+            e.list().unique_stable().into()
+        } else {
+            e.list().unique().into()
+        }
+    }
+}

data/ext/polars/src/{lazy → expr}/meta.rs

@@ -1,6 +1,12 @@
-use crate::{RArray, RbExpr, RbPolarsErr, RbResult};
+use magnus::RArray;
+
+use crate::{RbExpr, RbPolarsErr, RbResult};
 
 impl RbExpr {
+    pub fn meta_eq(&self, other: &RbExpr) -> bool {
+        self.inner == other.inner
+    }
+
     pub fn meta_pop(&self) -> RArray {
         RArray::from_iter(
             self.inner
@@ -12,11 +18,7 @@ impl RbExpr {
         )
     }
 
-    pub fn meta_eq(&self, other: &RbExpr) -> bool {
-        self.inner == other.inner
-    }
-
-    pub fn meta_roots(&self) -> Vec<String> {
+    pub fn meta_root_names(&self) -> Vec<String> {
         self.inner
             .clone()
             .meta()
@@ -39,4 +41,12 @@ impl RbExpr {
     pub fn meta_undo_aliases(&self) -> RbExpr {
         self.inner.clone().meta().undo_aliases().into()
     }
+
+    pub fn meta_has_multiple_outputs(&self) -> bool {
+        self.inner.clone().meta().has_multiple_outputs()
+    }
+
+    pub fn meta_is_regex_projection(&self) -> bool {
+        self.inner.clone().meta().is_regex_projection()
+    }
 }

data/ext/polars/src/expr/string.rs

@@ -0,0 +1,314 @@
+use polars::prelude::*;
+
+use crate::conversion::Wrap;
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn str_concat(&self, delimiter: String) -> Self {
+        self.inner.clone().str().concat(&delimiter).into()
+    }
+
+    pub fn str_to_date(
+        &self,
+        format: Option<String>,
+        strict: bool,
+        exact: bool,
+        cache: bool,
+    ) -> Self {
+        let options = StrptimeOptions {
+            format,
+            strict,
+            exact,
+            cache,
+        };
+        self.inner.clone().str().to_date(options).into()
+    }
+
+    pub fn str_to_datetime(
+        &self,
+        format: Option<String>,
+        time_unit: Option<Wrap<TimeUnit>>,
+        time_zone: Option<TimeZone>,
+        strict: bool,
+        exact: bool,
+        cache: bool,
+    ) -> Self {
+        let options = StrptimeOptions {
+            format,
+            strict,
+            exact,
+            cache,
+        };
+        self.inner
+            .clone()
+            .str()
+            .to_datetime(time_unit.map(|tu| tu.0), time_zone, options)
+            .into()
+    }
+
+    pub fn str_to_time(&self, format: Option<String>, strict: bool, cache: bool) -> Self {
+        let options = StrptimeOptions {
+            format,
+            strict,
+            cache,
+            exact: true,
+        };
+        self.inner.clone().str().to_time(options).into()
+    }
+
+    pub fn str_strip(&self, matches: Option<String>) -> Self {
+        self.inner.clone().str().strip(matches).into()
+    }
+
+    pub fn str_rstrip(&self, matches: Option<String>) -> Self {
+        self.inner.clone().str().rstrip(matches).into()
+    }
+
+    pub fn str_lstrip(&self, matches: Option<String>) -> Self {
+        self.inner.clone().str().lstrip(matches).into()
+    }
+
+    pub fn str_slice(&self, start: i64, length: Option<u64>) -> Self {
+        let function = move |s: Series| {
+            let ca = s.utf8()?;
+            Ok(Some(ca.str_slice(start, length)?.into_series()))
+        };
+        self.clone()
+            .inner
+            .map(function, GetOutput::from_type(DataType::Utf8))
+            .with_fmt("str.slice")
+            .into()
+    }
+
+    pub fn str_explode(&self) -> Self {
+        self.inner.clone().str().explode().into()
+    }
+
+    pub fn str_to_uppercase(&self) -> Self {
+        self.inner.clone().str().to_uppercase().into()
+    }
+
+    pub fn str_to_lowercase(&self) -> Self {
+        self.inner.clone().str().to_lowercase().into()
+    }
+
+    pub fn str_lengths(&self) -> Self {
+        let function = |s: Series| {
+            let ca = s.utf8()?;
+            Ok(Some(ca.str_lengths().into_series()))
+        };
+        self.clone()
+            .inner
+            .map(function, GetOutput::from_type(DataType::UInt32))
+            .with_fmt("str.lengths")
+            .into()
+    }
+
+    pub fn str_n_chars(&self) -> Self {
+        let function = |s: Series| {
+            let ca = s.utf8()?;
+            Ok(Some(ca.str_n_chars().into_series()))
+        };
+        self.clone()
+            .inner
+            .map(function, GetOutput::from_type(DataType::UInt32))
+            .with_fmt("str.n_chars")
+            .into()
+    }
+
+    pub fn str_replace_n(&self, pat: &RbExpr, val: &RbExpr, literal: bool, n: i64) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .replace_n(pat.inner.clone(), val.inner.clone(), literal, n)
+            .into()
+    }
+
+    pub fn str_replace_all(&self, pat: &RbExpr, val: &RbExpr, literal: bool) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .replace_all(pat.inner.clone(), val.inner.clone(), literal)
+            .into()
+    }
+
+    pub fn str_zfill(&self, alignment: usize) -> Self {
+        self.clone().inner.str().zfill(alignment).into()
+    }
+
+    pub fn str_ljust(&self, width: usize, fillchar: char) -> Self {
+        self.clone().inner.str().ljust(width, fillchar).into()
+    }
+
+    pub fn str_rjust(&self, width: usize, fillchar: char) -> Self {
+        self.clone().inner.str().rjust(width, fillchar).into()
+    }
+
+    pub fn str_contains(&self, pat: &RbExpr, literal: Option<bool>, strict: bool) -> Self {
+        match literal {
+            Some(true) => self
+                .inner
+                .clone()
+                .str()
+                .contains_literal(pat.inner.clone())
+                .into(),
+            _ => self
+                .inner
+                .clone()
+                .str()
+                .contains(pat.inner.clone(), strict)
+                .into(),
+        }
+    }
+
+    pub fn str_ends_with(&self, sub: &RbExpr) -> Self {
+        self.inner.clone().str().ends_with(sub.inner.clone()).into()
+    }
+
+    pub fn str_starts_with(&self, sub: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .starts_with(sub.inner.clone())
+            .into()
+    }
+
+    pub fn str_hex_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8().map(|s| Some(s.hex_encode().into_series())),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.hex_encode")
+            .into()
+    }
+
+    pub fn str_hex_decode(&self, strict: bool) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8()?.hex_decode(strict).map(|s| Some(s.into_series())),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.hex_decode")
+            .into()
+    }
+
+    pub fn str_base64_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.utf8().map(|s| Some(s.base64_encode().into_series())),
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.base64_encode")
+            .into()
+    }
+
+    pub fn str_base64_decode(&self, strict: bool) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| {
+                    s.utf8()?
+                        .base64_decode(strict)
+                        .map(|s| Some(s.into_series()))
+                },
+                GetOutput::same_type(),
+            )
+            .with_fmt("str.base64_decode")
+            .into()
+    }
+
+    pub fn str_parse_int(&self, radix: u32, strict: bool) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .from_radix(radix, strict)
+            .with_fmt("str.parse_int")
+            .into()
+    }
+
+    pub fn str_json_extract(
+        &self,
+        dtype: Option<Wrap<DataType>>,
+        infer_schema_len: Option<usize>,
+    ) -> Self {
+        let dtype = dtype.map(|wrap| wrap.0);
+
+        let output_type = match dtype.clone() {
+            Some(dtype) => GetOutput::from_type(dtype),
+            None => GetOutput::from_type(DataType::Unknown),
+        };
+
+        let function = move |s: Series| {
+            let ca = s.utf8()?;
+            match ca.json_extract(dtype.clone(), infer_schema_len) {
+                Ok(ca) => Ok(Some(ca.into_series())),
+                Err(e) => Err(PolarsError::ComputeError(format!("{e:?}").into())),
+            }
+        };
+
+        self.clone()
+            .inner
+            .map(function, output_type)
+            .with_fmt("str.json_extract")
+            .into()
+    }
+
+    pub fn str_json_path_match(&self, pat: String) -> Self {
+        let function = move |s: Series| {
+            let ca = s.utf8()?;
+            match ca.json_path_match(&pat) {
+                Ok(ca) => Ok(Some(ca.into_series())),
+                Err(e) => Err(PolarsError::ComputeError(format!("{:?}", e).into())),
+            }
+        };
+        self.clone()
+            .inner
+            .map(function, GetOutput::from_type(DataType::Utf8))
+            .with_fmt("str.json_path_match")
+            .into()
+    }
+
+    pub fn str_extract(&self, pat: String, group_index: usize) -> Self {
+        self.inner.clone().str().extract(&pat, group_index).into()
+    }
+
+    pub fn str_extract_all(&self, pat: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .extract_all(pat.inner.clone())
+            .into()
+    }
+
+    pub fn str_count_match(&self, pat: String) -> Self {
+        self.inner.clone().str().count_match(&pat).into()
+    }
+
+    pub fn str_split(&self, by: String) -> Self {
+        self.inner.clone().str().split(&by).into()
+    }
+
+    pub fn str_split_inclusive(&self, by: String) -> Self {
+        self.inner.clone().str().split_inclusive(&by).into()
+    }
+
+    pub fn str_split_exact(&self, by: String, n: usize) -> Self {
+        self.inner.clone().str().split_exact(&by, n).into()
+    }
+
+    pub fn str_split_exact_inclusive(&self, by: String, n: usize) -> Self {
+        self.inner
+            .clone()
+            .str()
+            .split_exact_inclusive(&by, n)
+            .into()
+    }
+
+    pub fn str_splitn(&self, by: String, n: usize) -> Self {
+        self.inner.clone().str().splitn(&by, n).into()
+    }
+}

data/ext/polars/src/expr/struct.rs

@@ -0,0 +1,15 @@
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn struct_field_by_name(&self, name: String) -> Self {
+        self.inner.clone().struct_().field_by_name(&name).into()
+    }
+
+    pub fn struct_field_by_index(&self, index: i64) -> Self {
+        self.inner.clone().struct_().field_by_index(index).into()
+    }
+
+    pub fn struct_rename_fields(&self, names: Vec<String>) -> Self {
+        self.inner.clone().struct_().rename_fields(names).into()
+    }
+}

data/ext/polars/src/expr.rs

@@ -0,0 +1,34 @@
+mod array;
+mod binary;
+mod categorical;
+mod datetime;
+mod general;
+mod list;
+mod meta;
+mod string;
+mod r#struct;
+
+use magnus::RArray;
+use polars::lazy::dsl::Expr;
+
+use crate::RbResult;
+
+#[magnus::wrap(class = "Polars::RbExpr")]
+#[derive(Clone)]
+pub struct RbExpr {
+    pub inner: Expr,
+}
+
+impl From<Expr> for RbExpr {
+    fn from(inner: Expr) -> Self {
+        RbExpr { inner }
+    }
+}
+
+pub fn rb_exprs_to_exprs(rb_exprs: RArray) -> RbResult<Vec<Expr>> {
+    let mut exprs = Vec::new();
+    for item in rb_exprs.each() {
+        exprs.push(item?.try_convert::<&RbExpr>()?.inner.clone());
+    }
+    Ok(exprs)
+}

data/ext/polars/src/functions/eager.rs

@@ -0,0 +1,93 @@
+use magnus::RArray;
+use polars::{functions, time};
+use polars_core::datatypes::{TimeUnit, TimeZone};
+use polars_core::prelude::{DataFrame, IntoSeries};
+
+use crate::conversion::{get_df, get_series, Wrap};
+use crate::error::RbPolarsErr;
+use crate::prelude::{ClosedWindow, Duration};
+use crate::{RbDataFrame, RbResult, RbSeries};
+
+pub fn concat_df(seq: RArray) -> RbResult<RbDataFrame> {
+    use polars_core::error::PolarsResult;
+
+    let mut iter = seq.each();
+    let first = iter.next().unwrap()?;
+
+    let first_rdf = get_df(first)?;
+    let identity_df = first_rdf.slice(0, 0);
+
+    let mut rdfs: Vec<PolarsResult<DataFrame>> = vec![Ok(first_rdf)];
+
+    for item in iter {
+        let rdf = get_df(item?)?;
+        rdfs.push(Ok(rdf));
+    }
+
+    let identity = Ok(identity_df);
+
+    let df = rdfs
+        .into_iter()
+        .fold(identity, |acc: PolarsResult<DataFrame>, df| {
+            let mut acc = acc?;
+            acc.vstack_mut(&df?)?;
+            Ok(acc)
+        })
+        .map_err(RbPolarsErr::from)?;
+
+    Ok(df.into())
+}
+
+pub fn concat_series(seq: RArray) -> RbResult<RbSeries> {
+    let mut iter = seq.each();
+    let first = iter.next().unwrap()?;
+
+    let mut s = get_series(first)?;
+
+    for res in iter {
+        let item = res?;
+        let item = get_series(item)?;
+        s.append(&item).map_err(RbPolarsErr::from)?;
+    }
+    Ok(s.into())
+}
+
+pub fn date_range(
+    start: i64,
+    stop: i64,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    name: String,
+    tu: Wrap<TimeUnit>,
+    tz: Option<TimeZone>,
+) -> RbResult<RbSeries> {
+    let date_range = time::date_range_impl(
+        &name,
+        start,
+        stop,
+        Duration::parse(&every),
+        closed.0,
+        tu.0,
+        tz.as_ref(),
+    )
+    .map_err(RbPolarsErr::from)?;
+    Ok(date_range.into_series().into())
+}
+
+pub fn diag_concat_df(seq: RArray) -> RbResult<RbDataFrame> {
+    let mut dfs = Vec::new();
+    for item in seq.each() {
+        dfs.push(get_df(item?)?);
+    }
+    let df = functions::diag_concat_df(&dfs).map_err(RbPolarsErr::from)?;
+    Ok(df.into())
+}
+
+pub fn hor_concat_df(seq: RArray) -> RbResult<RbDataFrame> {
+    let mut dfs = Vec::new();
+    for item in seq.each() {
+        dfs.push(get_df(item?)?);
+    }
+    let df = functions::hor_concat_df(&dfs).map_err(RbPolarsErr::from)?;
+    Ok(df.into())
+}

data/ext/polars/src/functions/io.rs

@@ -0,0 +1,34 @@
+use magnus::{RHash, Value};
+
+use crate::conversion::Wrap;
+use crate::file::get_file_like;
+use crate::prelude::DataType;
+use crate::{RbPolarsErr, RbResult};
+
+pub fn read_ipc_schema(rb_f: Value) -> RbResult<Value> {
+    use polars_core::export::arrow::io::ipc::read::read_file_metadata;
+    let mut r = get_file_like(rb_f, false)?;
+    let metadata = read_file_metadata(&mut r).map_err(RbPolarsErr::arrow)?;
+
+    let dict = RHash::new();
+    for field in metadata.schema.fields {
+        let dt: Wrap<DataType> = Wrap((&field.data_type).into());
+        dict.aset(field.name, dt)?;
+    }
+    Ok(dict.into())
+}
+
+pub fn read_parquet_schema(rb_f: Value) -> RbResult<Value> {
+    use polars_core::export::arrow::io::parquet::read::{infer_schema, read_metadata};
+
+    let mut r = get_file_like(rb_f, false)?;
+    let metadata = read_metadata(&mut r).map_err(RbPolarsErr::arrow)?;
+    let arrow_schema = infer_schema(&metadata).map_err(RbPolarsErr::arrow)?;
+
+    let dict = RHash::new();
+    for field in arrow_schema.fields {
+        let dt: Wrap<DataType> = Wrap((&field.data_type).into());
+        dict.aset(field.name, dt)?;
+    }
+    Ok(dict.into())
+}