polars-df 0.20.0 → 0.21.1

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

@@ -1,4 +1,4 @@
-use magnus::{r_hash::ForEach, typed_data::Obj, IntoValue, RArray, RHash, TryConvert, Value};
+use magnus::{IntoValue, RArray, RHash, TryConvert, Value, r_hash::ForEach, typed_data::Obj};
 use polars::io::{HiveOptions, RowIndex};
 use polars::lazy::frame::LazyFrame;
 use polars::prelude::*;
@@ -6,16 +6,17 @@ use polars_plan::dsl::ScanSources;
 use std::cell::RefCell;
 use std::io::BufWriter;
 use std::num::NonZeroUsize;
-use std::path::PathBuf;
 
 use super::SinkTarget;
 use crate::conversion::*;
 use crate::expr::rb_exprs_to_exprs;
+use crate::expr::selector::RbSelector;
 use crate::file::get_file_like;
+use crate::io::RbScanOptions;
 use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbLazyGroupBy, RbPolarsErr, RbResult, RbValueError};
 
-fn rbobject_to_first_path_and_scan_sources(obj: Value) -> RbResult<(Option<PathBuf>, ScanSources)> {
-    use crate::file::{get_ruby_scan_source_input, RubyScanSourceInput};
+fn rbobject_to_first_path_and_scan_sources(obj: Value) -> RbResult<(Option<PlPath>, ScanSources)> {
+    use crate::file::{RubyScanSourceInput, get_ruby_scan_source_input};
     Ok(match get_ruby_scan_source_input(obj, false)? {
         RubyScanSourceInput::Path(path) => (Some(path.clone()), ScanSources::Paths([path].into())),
         RubyScanSourceInput::File(file) => (None, ScanSources::Files([file].into())),
@@ -43,7 +44,7 @@ impl RbLazyFrame {
 
         let sources = sources.0;
         let (_first_path, sources) = match source {
-            None => (sources.first_path().map(|p| p.to_path_buf()), sources),
+            None => (sources.first_path().map(|p| p.into_owned()), sources),
             Some(source) => rbobject_to_first_path_and_scan_sources(source)?,
         };
 
@@ -111,7 +112,7 @@ impl RbLazyFrame {
 
         let sources = sources.0;
         let (_first_path, sources) = match source {
-            None => (sources.first_path().map(|p| p.to_path_buf()), sources),
+            None => (sources.first_path().map(|p| p.into_owned()), sources),
             Some(source) => rbobject_to_first_path_and_scan_sources(source)?,
         };
 
@@ -147,72 +148,35 @@ impl RbLazyFrame {
         Ok(r.finish().map_err(RbPolarsErr::from)?.into())
     }
 
-    pub fn new_from_parquet(arguments: &[Value]) -> RbResult<Self> {
-        let source = Option::<Value>::try_convert(arguments[0])?;
-        let sources = Wrap::<ScanSources>::try_convert(arguments[1])?;
-        let n_rows = Option::<usize>::try_convert(arguments[2])?;
-        let cache = bool::try_convert(arguments[3])?;
-        let parallel = Wrap::<ParallelStrategy>::try_convert(arguments[4])?;
-        let rechunk = bool::try_convert(arguments[5])?;
-        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[6])?;
-        let low_memory = bool::try_convert(arguments[7])?;
-        let cloud_options = Option::<Vec<(String, String)>>::try_convert(arguments[8])?;
-        let _credential_provider = Option::<Value>::try_convert(arguments[9])?;
-        let use_statistics = bool::try_convert(arguments[10])?;
-        let hive_partitioning = Option::<bool>::try_convert(arguments[11])?;
-        let schema = Option::<Wrap<Schema>>::try_convert(arguments[12])?;
-        let hive_schema = Option::<Wrap<Schema>>::try_convert(arguments[13])?;
-        let try_parse_hive_dates = bool::try_convert(arguments[14])?;
-        let retries = usize::try_convert(arguments[15])?;
-        let glob = bool::try_convert(arguments[16])?;
-        let include_file_paths = Option::<String>::try_convert(arguments[17])?;
-        let allow_missing_columns = bool::try_convert(arguments[18])?;
+    pub fn new_from_parquet(
+        sources: Wrap<ScanSources>,
+        schema: Option<Wrap<Schema>>,
+        scan_options: RbScanOptions,
+        parallel: Wrap<ParallelStrategy>,
+        low_memory: bool,
+        use_statistics: bool,
+    ) -> RbResult<Self> {
+        use crate::utils::to_rb_err;
 
         let parallel = parallel.0;
-        let hive_schema = hive_schema.map(|s| Arc::new(s.0));
-
-        let row_index = row_index.map(|(name, offset)| RowIndex {
-            name: name.into(),
-            offset,
-        });
-
-        let hive_options = HiveOptions {
-            enabled: hive_partitioning,
-            hive_start_idx: 0,
-            schema: hive_schema,
-            try_parse_dates: try_parse_hive_dates,
-        };
 
-        let mut args = ScanArgsParquet {
-            n_rows,
-            cache,
+        let options = ParquetOptions {
+            schema: schema.map(|x| Arc::new(x.0)),
             parallel,
-            rechunk,
-            row_index,
             low_memory,
-            cloud_options: None,
             use_statistics,
-            schema: schema.map(|x| Arc::new(x.0)),
-            hive_options,
-            glob,
-            include_file_paths: include_file_paths.map(|x| x.into()),
-            allow_missing_columns,
         };
 
         let sources = sources.0;
-        let (first_path, sources) = match source {
-            None => (sources.first_path().map(|p| p.to_path_buf()), sources),
-            Some(source) => rbobject_to_first_path_and_scan_sources(source)?,
-        };
+        let first_path = sources.first_path().map(|p| p.into_owned());
 
-        if let Some(first_path) = first_path {
-            let first_path_url = first_path.to_string_lossy();
-            let cloud_options =
-                parse_cloud_options(&first_path_url, cloud_options.unwrap_or_default())?;
-            args.cloud_options = Some(cloud_options.with_max_retries(retries));
-        }
+        let unified_scan_args =
+            scan_options.extract_unified_scan_args(first_path.as_ref().map(|p| p.as_ref()))?;
 
-        let lf = LazyFrame::scan_parquet_sources(sources, args).map_err(RbPolarsErr::from)?;
+        let lf: LazyFrame = DslBuilder::scan_parquet(sources, options, unified_scan_args)
+            .map_err(to_rb_err)?
+            .build()
+            .into();
 
         Ok(lf.into())
     }
@@ -254,7 +218,7 @@ impl RbLazyFrame {
 
         let sources = sources.0;
         let (_first_path, sources) = match source {
-            None => (sources.first_path().map(|p| p.to_path_buf()), sources),
+            None => (sources.first_path().map(|p| p.into_owned()), sources),
             Some(source) => rbobject_to_first_path_and_scan_sources(source)?,
         };
 
@@ -265,7 +229,7 @@ impl RbLazyFrame {
     pub fn write_json(&self, rb_f: Value) -> RbResult<()> {
         let file = BufWriter::new(get_file_like(rb_f, true)?);
         serde_json::to_writer(file, &self.ldf.borrow().logical_plan)
-            .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
+            .map_err(|err| RbValueError::new_err(format!("{err:?}")))?;
         Ok(())
     }
 
@@ -361,6 +325,30 @@ impl RbLazyFrame {
             .into())
     }
 
+    pub fn top_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
+        let ldf = self.ldf.borrow().clone();
+        let exprs = rb_exprs_to_exprs(by)?;
+        Ok(ldf
+            .top_k(
+                k,
+                exprs,
+                SortMultipleOptions::new().with_order_descending_multi(reverse),
+            )
+            .into())
+    }
+
+    pub fn bottom_k(&self, k: IdxSize, by: RArray, reverse: Vec<bool>) -> RbResult<Self> {
+        let ldf = self.ldf.borrow().clone();
+        let exprs = rb_exprs_to_exprs(by)?;
+        Ok(ldf
+            .bottom_k(
+                k,
+                exprs,
+                SortMultipleOptions::new().with_order_descending_multi(reverse),
+            )
+            .into())
+    }
+
     pub fn cache(&self) -> Self {
         let ldf = self.ldf.borrow().clone();
         ldf.cache().into()
@@ -399,10 +387,8 @@ impl RbLazyFrame {
         let cloud_options = match target.base_path() {
             None => None,
             Some(base_path) => {
-                let cloud_options = parse_cloud_options(
-                    base_path.to_str().unwrap(),
-                    cloud_options.unwrap_or_default(),
-                )?;
+                let cloud_options =
+                    parse_cloud_options(base_path.to_str(), cloud_options.unwrap_or_default())?;
                 Some(cloud_options.with_max_retries(retries))
             }
         };
@@ -434,10 +420,8 @@ impl RbLazyFrame {
         let cloud_options = match target.base_path() {
             None => None,
             Some(base_path) => {
-                let cloud_options = parse_cloud_options(
-                    base_path.to_str().unwrap(),
-                    cloud_options.unwrap_or_default(),
-                )?;
+                let cloud_options =
+                    parse_cloud_options(base_path.to_str(), cloud_options.unwrap_or_default())?;
                 Some(cloud_options.with_max_retries(retries))
             }
         };
@@ -466,11 +450,12 @@ impl RbLazyFrame {
         let time_format = Option::<String>::try_convert(arguments[9])?;
         let float_scientific = Option::<bool>::try_convert(arguments[10])?;
         let float_precision = Option::<usize>::try_convert(arguments[11])?;
-        let null_value = Option::<String>::try_convert(arguments[12])?;
-        let quote_style = Option::<Wrap<QuoteStyle>>::try_convert(arguments[13])?;
-        let cloud_options = Option::<Vec<(String, String)>>::try_convert(arguments[14])?;
-        let retries = usize::try_convert(arguments[15])?;
-        let sink_options = Wrap::<SinkOptions>::try_convert(arguments[16])?;
+        let decimal_comma = bool::try_convert(arguments[12])?;
+        let null_value = Option::<String>::try_convert(arguments[13])?;
+        let quote_style = Option::<Wrap<QuoteStyle>>::try_convert(arguments[14])?;
+        let cloud_options = Option::<Vec<(String, String)>>::try_convert(arguments[15])?;
+        let retries = usize::try_convert(arguments[16])?;
+        let sink_options = Wrap::<SinkOptions>::try_convert(arguments[17])?;
 
         let quote_style = quote_style.map_or(QuoteStyle::default(), |wrap| wrap.0);
         let null_value = null_value.unwrap_or(SerializeOptions::default().null);
@@ -481,6 +466,7 @@ impl RbLazyFrame {
             datetime_format,
             float_scientific,
             float_precision,
+            decimal_comma,
             separator,
             quote_char,
             null: null_value,
@@ -498,10 +484,8 @@ impl RbLazyFrame {
         let cloud_options = match target.base_path() {
             None => None,
             Some(base_path) => {
-                let cloud_options = parse_cloud_options(
-                    base_path.to_str().unwrap(),
-                    cloud_options.unwrap_or_default(),
-                )?;
+                let cloud_options =
+                    parse_cloud_options(base_path.to_str(), cloud_options.unwrap_or_default())?;
                 Some(cloud_options.with_max_retries(retries))
             }
         };
@@ -529,10 +513,8 @@ impl RbLazyFrame {
         let cloud_options = match target.base_path() {
             None => None,
             Some(base_path) => {
-                let cloud_options = parse_cloud_options(
-                    base_path.to_str().unwrap(),
-                    cloud_options.unwrap_or_default(),
-                )?;
+                let cloud_options =
+                    parse_cloud_options(base_path.to_str(), cloud_options.unwrap_or_default())?;
                 Some(cloud_options.with_max_retries(retries))
             }
         };
@@ -546,15 +528,14 @@ impl RbLazyFrame {
         .map_err(Into::into)
     }
 
-    pub fn fetch(&self, n_rows: usize) -> RbResult<RbDataFrame> {
+    pub fn filter(&self, predicate: &RbExpr) -> Self {
         let ldf = self.ldf.borrow().clone();
-        let df = ldf.fetch(n_rows).map_err(RbPolarsErr::from)?;
-        Ok(df.into())
+        ldf.filter(predicate.inner.clone()).into()
     }
 
-    pub fn filter(&self, predicate: &RbExpr) -> Self {
+    pub fn remove(&self, predicate: &RbExpr) -> Self {
         let ldf = self.ldf.borrow().clone();
-        ldf.filter(predicate.inner.clone()).into()
+        ldf.remove(predicate.inner.clone()).into()
     }
 
     pub fn select(&self, exprs: RArray) -> RbResult<Self> {
@@ -689,15 +670,19 @@ impl RbLazyFrame {
             .allow_parallel(allow_parallel)
             .force_parallel(force_parallel)
             .coalesce(coalesce)
-            .how(JoinType::AsOf(AsOfOptions {
+            .how(JoinType::AsOf(Box::new(AsOfOptions {
                 strategy: strategy.0,
                 left_by: left_by.map(strings_to_pl_smallstr),
                 right_by: right_by.map(strings_to_pl_smallstr),
-                tolerance: tolerance.map(|t| t.0.into_static()),
+                tolerance: tolerance.map(|t| {
+                    let av = t.0.into_static();
+                    let dtype = av.dtype();
+                    Scalar::new(dtype, av)
+                }),
                 tolerance_str: tolerance_str.map(|s| s.into()),
                 allow_eq,
                 check_sortedness,
-            }))
+            })))
             .suffix(suffix)
             .finish()
             .into())
@@ -745,6 +730,20 @@ impl RbLazyFrame {
             .into())
     }
 
+    pub fn join_where(&self, other: &Self, predicates: RArray, suffix: String) -> RbResult<Self> {
+        let ldf = self.ldf.borrow().clone();
+        let other = other.ldf.borrow().clone();
+
+        let predicates = rb_exprs_to_exprs(predicates)?;
+
+        Ok(ldf
+            .join_builder()
+            .with(other)
+            .suffix(suffix)
+            .join_where(predicates)
+            .into())
+    }
+
     pub fn with_column(&self, expr: &RbExpr) -> Self {
         let ldf = self.ldf.borrow().clone();
         ldf.with_column(expr.inner.clone()).into()
@@ -832,10 +831,12 @@ impl RbLazyFrame {
         out.into()
     }
 
-    pub fn explode(&self, column: RArray) -> RbResult<Self> {
-        let ldf = self.ldf.borrow().clone();
-        let column = rb_exprs_to_exprs(column)?;
-        Ok(ldf.explode(column).into())
+    pub fn explode(&self, subset: &RbSelector) -> Self {
+        self.ldf
+            .borrow()
+            .clone()
+            .explode(subset.inner.clone())
+            .into()
     }
 
     pub fn null_count(&self) -> Self {
@@ -846,10 +847,11 @@ impl RbLazyFrame {
     pub fn unique(
         &self,
         maintain_order: bool,
-        subset: Option<Vec<String>>,
+        subset: Option<&RbSelector>,
         keep: Wrap<UniqueKeepStrategy>,
     ) -> RbResult<Self> {
         let ldf = self.ldf.borrow().clone();
+        let subset = subset.map(|e| e.inner.clone());
         Ok(match maintain_order {
             true => ldf.unique_stable_generic(subset, keep.0),
             false => ldf.unique_generic(subset, keep.0),
@@ -857,9 +859,19 @@ impl RbLazyFrame {
         .into())
     }
 
-    pub fn drop_nulls(&self, subset: Option<Vec<String>>) -> Self {
-        let ldf = self.ldf.borrow().clone();
-        ldf.drop_nulls(subset.map(|v| v.into_iter().map(|s| col(&s)).collect()))
+    pub fn drop_nans(&self, subset: Option<&RbSelector>) -> Self {
+        self.ldf
+            .borrow()
+            .clone()
+            .drop_nans(subset.map(|e| e.inner.clone()))
+            .into()
+    }
+
+    pub fn drop_nulls(&self, subset: Option<&RbSelector>) -> Self {
+        self.ldf
+            .borrow()
+            .clone()
+            .drop_nulls(subset.map(|e| e.inner.clone()))
             .into()
     }
 
@@ -875,16 +887,14 @@ impl RbLazyFrame {
 
     pub fn unpivot(
         &self,
-        on: RArray,
-        index: RArray,
+        on: &RbSelector,
+        index: &RbSelector,
         value_name: Option<String>,
         variable_name: Option<String>,
     ) -> RbResult<Self> {
-        let on = rb_exprs_to_exprs(on)?;
-        let index = rb_exprs_to_exprs(index)?;
         let args = UnpivotArgsDSL {
-            on: on.into_iter().map(|e| e.into()).collect(),
-            index: index.into_iter().map(|e| e.into()).collect(),
+            on: on.inner.clone(),
+            index: index.inner.clone(),
             value_name: value_name.map(|s| s.into()),
             variable_name: variable_name.map(|s| s.into()),
         };
@@ -898,9 +908,8 @@ impl RbLazyFrame {
         ldf.with_row_index(&name, offset).into()
     }
 
-    pub fn drop(&self, cols: Vec<String>) -> Self {
-        let ldf = self.ldf.borrow().clone();
-        ldf.drop(cols).into()
+    pub fn drop(&self, columns: &RbSelector) -> Self {
+        self.ldf.borrow().clone().drop(columns.inner.clone()).into()
     }
 
     pub fn cast(&self, rb_dtypes: RHash, strict: bool) -> RbResult<Self> {
@@ -941,8 +950,12 @@ impl RbLazyFrame {
         Ok(schema_dict)
     }
 
-    pub fn unnest(&self, cols: Vec<String>) -> Self {
-        self.ldf.borrow().clone().unnest(cols).into()
+    pub fn unnest(&self, columns: &RbSelector) -> Self {
+        self.ldf
+            .borrow()
+            .clone()
+            .unnest(columns.inner.clone())
+            .into()
     }
 
     pub fn count(&self) -> Self {

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

@@ -25,7 +25,7 @@ impl RbLazyFrame {
         let json = unsafe { std::mem::transmute::<&'_ str, &'static str>(json.as_str()) };
 
         let lp = serde_json::from_str::<DslPlan>(json)
-            .map_err(|err| RbValueError::new_err(format!("{:?}", err)))?;
+            .map_err(|err| RbValueError::new_err(format!("{err:?}")))?;
         Ok(LazyFrame::from(lp).into())
     }
 }

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

@@ -1,9 +1,9 @@
-use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex};
 
 use magnus::{RHash, TryConvert, Value};
 use polars::prelude::sync_on_close::SyncOnCloseType;
-use polars::prelude::{SinkOptions, SpecialEq};
+use polars::prelude::{PlPath, SinkOptions, SpecialEq};
+use polars_utils::plpath::PlPathRef;
 
 use crate::prelude::Wrap;
 use crate::{RbResult, RbValueError};
@@ -15,8 +15,8 @@ pub enum SinkTarget {
 
 impl TryConvert for Wrap<polars_plan::dsl::SinkTarget> {
     fn try_convert(ob: Value) -> RbResult<Self> {
-        if let Ok(v) = PathBuf::try_convert(ob) {
-            Ok(Wrap(polars::prelude::SinkTarget::Path(Arc::new(v))))
+        if let Ok(v) = String::try_convert(ob) {
+            Ok(Wrap(polars::prelude::SinkTarget::Path(PlPath::new(&v))))
         } else {
             let writer = {
                 let rb_f = ob;
@@ -39,10 +39,10 @@ impl TryConvert for SinkTarget {
 }
 
 impl SinkTarget {
-    pub fn base_path(&self) -> Option<&Path> {
+    pub fn base_path(&self) -> Option<PlPathRef<'_>> {
         match self {
             Self::File(t) => match t {
-                polars::prelude::SinkTarget::Path(p) => Some(p.as_path()),
+                polars::prelude::SinkTarget::Path(p) => Some(p.as_ref()),
                 polars::prelude::SinkTarget::Dyn(_) => None,
             },
         }