polars-df 0.9.0 → 0.11.0

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

@@ -39,3 +39,77 @@ pub fn date_range(
     let time_unit = time_unit.map(|x| x.0);
     dsl::date_range(start, end, every, closed, time_unit, time_zone).into()
 }
+
+pub fn date_ranges(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    time_unit: Option<Wrap<TimeUnit>>,
+    time_zone: Option<TimeZone>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    let time_unit = time_unit.map(|x| x.0);
+    dsl::date_ranges(start, end, every, closed, time_unit, time_zone).into()
+}
+
+pub fn datetime_range(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    time_unit: Option<Wrap<TimeUnit>>,
+    time_zone: Option<TimeZone>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    let time_unit = time_unit.map(|x| x.0);
+    dsl::datetime_range(start, end, every, closed, time_unit, time_zone).into()
+}
+
+pub fn datetime_ranges(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    time_unit: Option<Wrap<TimeUnit>>,
+    time_zone: Option<TimeZone>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    let time_unit = time_unit.map(|x| x.0);
+    dsl::datetime_ranges(start, end, every, closed, time_unit, time_zone).into()
+}
+
+pub fn time_range(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    dsl::time_range(start, end, every, closed).into()
+}
+
+pub fn time_ranges(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    dsl::time_ranges(start, end, every, closed).into()
+}

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

@@ -2,42 +2,72 @@ use polars::lazy::dsl;
 
 use crate::RbExpr;
 
+pub fn when(condition: &RbExpr) -> RbWhen {
+    RbWhen {
+        inner: dsl::when(condition.inner.clone()),
+    }
+}
+
 #[magnus::wrap(class = "Polars::RbWhen")]
 #[derive(Clone)]
 pub struct RbWhen {
     pub inner: dsl::When,
 }
 
-impl From<dsl::When> for RbWhen {
-    fn from(inner: dsl::When) -> Self {
-        RbWhen { inner }
-    }
-}
-
-#[magnus::wrap(class = "Polars::RbWhenThen")]
+#[magnus::wrap(class = "Polars::RbThen")]
 #[derive(Clone)]
 pub struct RbThen {
     pub inner: dsl::Then,
 }
 
-impl From<dsl::Then> for RbThen {
-    fn from(inner: dsl::Then) -> Self {
-        RbThen { inner }
-    }
+#[magnus::wrap(class = "Polars::RbChainedWhen")]
+#[derive(Clone)]
+pub struct RbChainedWhen {
+    pub inner: dsl::ChainedWhen,
+}
+
+#[magnus::wrap(class = "Polars::RbChainedThen")]
+#[derive(Clone)]
+pub struct RbChainedThen {
+    pub inner: dsl::ChainedThen,
 }
 
 impl RbWhen {
-    pub fn then(&self, expr: &RbExpr) -> RbThen {
-        self.inner.clone().then(expr.inner.clone()).into()
+    pub fn then(&self, statement: &RbExpr) -> RbThen {
+        RbThen {
+            inner: self.inner.clone().then(statement.inner.clone()),
+        }
     }
 }
 
 impl RbThen {
-    pub fn overwise(&self, expr: &RbExpr) -> RbExpr {
-        self.inner.clone().otherwise(expr.inner.clone()).into()
+    pub fn when(&self, condition: &RbExpr) -> RbChainedWhen {
+        RbChainedWhen {
+            inner: self.inner.clone().when(condition.inner.clone()),
+        }
+    }
+
+    pub fn otherwise(&self, statement: &RbExpr) -> RbExpr {
+        self.inner.clone().otherwise(statement.inner.clone()).into()
     }
 }
 
-pub fn when(predicate: &RbExpr) -> RbWhen {
-    dsl::when(predicate.inner.clone()).into()
+impl RbChainedWhen {
+    pub fn then(&self, statement: &RbExpr) -> RbChainedThen {
+        RbChainedThen {
+            inner: self.inner.clone().then(statement.inner.clone()),
+        }
+    }
+}
+
+impl RbChainedThen {
+    pub fn when(&self, condition: &RbExpr) -> RbChainedWhen {
+        RbChainedWhen {
+            inner: self.inner.clone().when(condition.inner.clone()),
+        }
+    }
+
+    pub fn otherwise(&self, statement: &RbExpr) -> RbExpr {
+        self.inner.clone().otherwise(statement.inner.clone()).into()
+    }
 }

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

@@ -0,0 +1 @@
+pub mod numo;

data/ext/polars/src/interop/numo/mod.rs

@@ -0,0 +1,2 @@
+pub mod to_numo_df;
+pub mod to_numo_series;

data/ext/polars/src/interop/numo/to_numo_df.rs

@@ -0,0 +1,23 @@
+use magnus::Value;
+use polars_core::utils::try_get_supertype;
+
+use crate::dataframe::RbDataFrame;
+
+impl RbDataFrame {
+    pub fn to_numo(&self) -> Option<Value> {
+        let mut st = None;
+        for s in self.df.borrow().iter() {
+            let dt_i = s.dtype();
+            match st {
+                None => st = Some(dt_i.clone()),
+                Some(ref mut st) => {
+                    *st = try_get_supertype(st, dt_i).ok()?;
+                }
+            }
+        }
+        let _st = st?;
+
+        // TODO
+        None
+    }
+}

data/ext/polars/src/interop/numo/to_numo_series.rs

@@ -0,0 +1,60 @@
+use magnus::{class, prelude::*, Module, RArray, RClass, RModule, Value};
+use polars_core::prelude::*;
+
+use crate::error::RbPolarsErr;
+use crate::raise_err;
+use crate::series::RbSeries;
+use crate::RbResult;
+
+impl RbSeries {
+    /// For numeric types, this should only be called for Series with null types.
+    /// This will cast to floats so that `nil = NAN`
+    pub fn to_numo(&self) -> RbResult<Value> {
+        let s = &self.series.borrow();
+        match s.dtype() {
+            DataType::String => {
+                let ca = s.str().unwrap();
+
+                // TODO make more efficient
+                let np_arr = RArray::from_iter(ca);
+                class::object()
+                    .const_get::<_, RModule>("Numo")?
+                    .const_get::<_, RClass>("RObject")?
+                    .funcall("cast", (np_arr,))
+            }
+            dt if dt.is_numeric() => {
+                if s.bit_repr_is_large() {
+                    let s = s.cast(&DataType::Float64).unwrap();
+                    let ca = s.f64().unwrap();
+                    // TODO make more efficient
+                    let np_arr = RArray::from_iter(ca.into_iter().map(|opt_v| match opt_v {
+                        Some(v) => v,
+                        None => f64::NAN,
+                    }));
+                    class::object()
+                        .const_get::<_, RModule>("Numo")?
+                        .const_get::<_, RClass>("DFloat")?
+                        .funcall("cast", (np_arr,))
+                } else {
+                    let s = s.cast(&DataType::Float32).unwrap();
+                    let ca = s.f32().unwrap();
+                    // TODO make more efficient
+                    let np_arr = RArray::from_iter(ca.into_iter().map(|opt_v| match opt_v {
+                        Some(v) => v,
+                        None => f32::NAN,
+                    }));
+                    class::object()
+                        .const_get::<_, RModule>("Numo")?
+                        .const_get::<_, RClass>("SFloat")?
+                        .funcall("cast", (np_arr,))
+                }
+            }
+            dt => {
+                raise_err!(
+                    format!("'to_numo' not supported for dtype: {dt:?}"),
+                    ComputeError
+                );
+            }
+        }
+    }
+}

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

@@ -1,5 +1,5 @@
 use magnus::{IntoValue, RArray, RHash, TryConvert, Value};
-use polars::io::RowIndex;
+use polars::io::{HiveOptions, RowIndex};
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
 use std::cell::RefCell;
@@ -48,7 +48,7 @@ impl RbLazyFrame {
         // in this scope
         let json = unsafe { std::mem::transmute::<&'_ str, &'static str>(json.as_str()) };
 
-        let lp = serde_json::from_str::<LogicalPlan>(json)
+        let lp = serde_json::from_str::<DslPlan>(json)
             .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
         Ok(LazyFrame::from(lp).into())
     }
@@ -63,7 +63,10 @@ impl RbLazyFrame {
         row_index: Option<(String, IdxSize)>,
     ) -> RbResult<Self> {
         let batch_size = batch_size.map(|v| v.0);
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
 
         let lf = LazyJsonLineReader::new(path)
             .with_infer_schema_length(infer_schema_length)
@@ -100,14 +103,17 @@ impl RbLazyFrame {
         let row_index = Option::<(String, IdxSize)>::try_convert(arguments[17])?;
         let try_parse_dates = bool::try_convert(arguments[18])?;
         let eol_char = String::try_convert(arguments[19])?;
+        let truncate_ragged_lines = bool::try_convert(arguments[20])?;
         // end arguments
 
         let null_values = null_values.map(|w| w.0);
         let quote_char = quote_char.map(|s| s.as_bytes()[0]);
         let separator = separator.as_bytes()[0];
         let eol_char = eol_char.as_bytes()[0];
-
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
 
         let overwrite_dtype = overwrite_dtype.map(|overwrite_dtype| {
             overwrite_dtype
@@ -115,25 +121,29 @@ impl RbLazyFrame {
                 .map(|(name, dtype)| Field::new(&name, dtype.0))
                 .collect::<Schema>()
         });
+
         let r = LazyCsvReader::new(path)
             .with_infer_schema_length(infer_schema_length)
             .with_separator(separator)
-            .has_header(has_header)
+            .with_has_header(has_header)
             .with_ignore_errors(ignore_errors)
             .with_skip_rows(skip_rows)
             .with_n_rows(n_rows)
             .with_cache(cache)
-            .with_dtype_overwrite(overwrite_dtype.as_ref())
-            .low_memory(low_memory)
+            .with_dtype_overwrite(overwrite_dtype.map(Arc::new))
+            // TODO add with_schema
+            .with_low_memory(low_memory)
             .with_comment_prefix(comment_prefix.as_deref())
             .with_quote_char(quote_char)
-            .with_end_of_line_char(eol_char)
+            .with_eol_char(eol_char)
             .with_rechunk(rechunk)
             .with_skip_rows_after_header(skip_rows_after_header)
             .with_encoding(encoding.0)
             .with_row_index(row_index)
             .with_try_parse_dates(try_parse_dates)
-            .with_null_values(null_values);
+            .with_null_values(null_values)
+            // TODO add with_missing_is_null
+            .with_truncate_ragged_lines(truncate_ragged_lines);
 
         if let Some(_lambda) = with_schema_modify {
             todo!();
@@ -144,7 +154,8 @@ impl RbLazyFrame {
 
     #[allow(clippy::too_many_arguments)]
     pub fn new_from_parquet(
-        path: String,
+        path: Option<PathBuf>,
+        paths: Vec<PathBuf>,
         n_rows: Option<usize>,
         cache: bool,
         parallel: Wrap<ParallelStrategy>,
@@ -153,21 +164,48 @@ impl RbLazyFrame {
         low_memory: bool,
         use_statistics: bool,
         hive_partitioning: bool,
+        hive_schema: Option<Wrap<Schema>>,
+        glob: bool,
     ) -> RbResult<Self> {
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let parallel = parallel.0;
+        let hive_schema = hive_schema.map(|s| Arc::new(s.0));
+
+        let first_path = if let Some(path) = &path {
+            path
+        } else {
+            paths
+                .first()
+                .ok_or_else(|| RbValueError::new_err("expected a path argument".to_string()))?
+        };
+
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
+        let hive_options = HiveOptions {
+            enabled: hive_partitioning,
+            schema: hive_schema,
+        };
+
         let args = ScanArgsParquet {
             n_rows,
             cache,
-            parallel: parallel.0,
+            parallel,
             rechunk,
             row_index,
             low_memory,
-            // TODO support cloud options
             cloud_options: None,
             use_statistics,
-            hive_partitioning,
+            hive_options,
+            glob,
         };
-        let lf = LazyFrame::scan_parquet(path, args).map_err(RbPolarsErr::from)?;
+
+        let lf = if path.is_some() {
+            LazyFrame::scan_parquet(first_path, args)
+        } else {
+            LazyFrame::scan_parquet_files(Arc::from(paths), args)
+        }
+        .map_err(RbPolarsErr::from)?;
         Ok(lf.into())
     }
 
@@ -179,13 +217,18 @@ impl RbLazyFrame {
         row_index: Option<(String, IdxSize)>,
         memory_map: bool,
     ) -> RbResult<Self> {
-        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: Arc::from(name.as_str()),
+            offset,
+        });
+
         let args = ScanArgsIpc {
             n_rows,
             cache,
             rechunk,
             row_index,
-            memmap: memory_map,
+            memory_map,
+            cloud_options: None,
         };
         let lf = LazyFrame::scan_ipc(path, args).map_err(RbPolarsErr::from)?;
         Ok(lf.into())
@@ -198,8 +241,11 @@ impl RbLazyFrame {
         Ok(())
     }
 
-    pub fn describe_plan(&self) -> String {
-        self.ldf.describe_plan()
+    pub fn describe_plan(&self) -> RbResult<String> {
+        self.ldf
+            .describe_plan()
+            .map_err(RbPolarsErr::from)
+            .map_err(Into::into)
     }
 
     pub fn describe_optimized_plan(&self) -> RbResult<String> {
@@ -242,17 +288,18 @@ impl RbLazyFrame {
     pub fn sort(
         &self,
         by_column: String,
-        reverse: bool,
+        descending: bool,
         nulls_last: bool,
         maintain_order: bool,
+        multithreaded: bool,
     ) -> Self {
         let ldf = self.ldf.clone();
         ldf.sort(
-            &by_column,
-            SortOptions {
-                descending: reverse,
+            [&by_column],
+            SortMultipleOptions {
+                descending: vec![descending],
                 nulls_last,
-                multithreaded: true,
+                multithreaded,
                 maintain_order,
             },
         )
@@ -261,15 +308,24 @@ impl RbLazyFrame {
 
     pub fn sort_by_exprs(
         &self,
-        by_column: RArray,
-        reverse: Vec<bool>,
+        by: RArray,
+        descending: Vec<bool>,
         nulls_last: bool,
         maintain_order: bool,
+        multithreaded: bool,
     ) -> RbResult<Self> {
         let ldf = self.ldf.clone();
-        let exprs = rb_exprs_to_exprs(by_column)?;
+        let exprs = rb_exprs_to_exprs(by)?;
         Ok(ldf
-            .sort_by_exprs(exprs, reverse, nulls_last, maintain_order)
+            .sort_by_exprs(
+                exprs,
+                SortMultipleOptions {
+                    descending,
+                    nulls_last,
+                    maintain_order,
+                    multithreaded,
+                },
+            )
             .into())
     }
 
@@ -326,6 +382,7 @@ impl RbLazyFrame {
         Ok(())
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn sink_csv(
         &self,
         path: PathBuf,
@@ -427,7 +484,7 @@ impl RbLazyFrame {
         let closed_window = closed.0;
         let ldf = self.ldf.clone();
         let by = rb_exprs_to_exprs(by)?;
-        let lazy_gb = ldf.group_by_rolling(
+        let lazy_gb = ldf.rolling(
             index_column.inner.clone(),
             by,
             RollingGroupOptions {
@@ -602,58 +659,56 @@ impl RbLazyFrame {
         ldf.fill_nan(fill_value.inner.clone()).into()
     }
 
-    pub fn min(&self) -> RbResult<Self> {
+    pub fn min(&self) -> Self {
         let ldf = self.ldf.clone();
-        let out = ldf.min().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+        let out = ldf.min();
+        out.into()
     }
 
-    pub fn max(&self) -> RbResult<Self> {
+    pub fn max(&self) -> Self {
         let ldf = self.ldf.clone();
-        let out = ldf.max().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+        let out = ldf.max();
+        out.into()
     }
 
-    pub fn sum(&self) -> RbResult<Self> {
+    pub fn sum(&self) -> Self {
         let ldf = self.ldf.clone();
-        let out = ldf.sum().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+        let out = ldf.sum();
+        out.into()
     }
 
-    pub fn mean(&self) -> RbResult<Self> {
+    pub fn mean(&self) -> Self {
         let ldf = self.ldf.clone();
-        let out = ldf.mean().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+        let out = ldf.mean();
+        out.into()
     }
 
-    pub fn std(&self, ddof: u8) -> RbResult<Self> {
+    pub fn std(&self, ddof: u8) -> Self {
         let ldf = self.ldf.clone();
-        let out = ldf.std(ddof).map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+        let out = ldf.std(ddof);
+        out.into()
     }
 
-    pub fn var(&self, ddof: u8) -> RbResult<Self> {
+    pub fn var(&self, ddof: u8) -> Self {
         let ldf = self.ldf.clone();
-        let out = ldf.var(ddof).map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+        let out = ldf.var(ddof);
+        out.into()
     }
 
-    pub fn median(&self) -> RbResult<Self> {
+    pub fn median(&self) -> Self {
         let ldf = self.ldf.clone();
-        let out = ldf.median().map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+        let out = ldf.median();
+        out.into()
     }
 
     pub fn quantile(
         &self,
         quantile: &RbExpr,
         interpolation: Wrap<QuantileInterpolOptions>,
-    ) -> RbResult<Self> {
+    ) -> Self {
         let ldf = self.ldf.clone();
-        let out = ldf
-            .quantile(quantile.inner.clone(), interpolation.0)
-            .map_err(RbPolarsErr::from)?;
-        Ok(out.into())
+        let out = ldf.quantile(quantile.inner.clone(), interpolation.0);
+        out.into()
     }
 
     pub fn explode(&self, column: RArray) -> RbResult<Self> {