polars-df 0.20.0 → 0.21.1

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

@@ -1,13 +1,13 @@
 use magnus::encoding::{self, EncodingCapable};
 use magnus::{
-    class, prelude::*, typed_data::Obj, value::Opaque, Float, Integer, RArray, RString, Ruby, Value,
+    Float, Integer, RArray, RString, Ruby, Value, class, prelude::*, typed_data::Obj, value::Opaque,
 };
 use polars::lazy::dsl;
 use polars::prelude::*;
 
-use crate::conversion::{get_lf, get_rbseq, Wrap};
+use crate::conversion::{Wrap, get_lf, get_rbseq};
+use crate::expr::datatype::RbDataTypeExpr;
 use crate::map::lazy::binary_lambda;
-use crate::prelude::vec_extract_wrapped;
 use crate::rb_exprs_to_exprs;
 use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
@@ -85,6 +85,10 @@ pub fn as_struct(exprs: RArray) -> RbResult<RbExpr> {
     Ok(dsl::as_struct(exprs).into())
 }
 
+pub fn field(names: Vec<String>) -> RbExpr {
+    dsl::Expr::Field(names.into_iter().map(|x| x.into()).collect()).into()
+}
+
 pub fn coalesce(exprs: RArray) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
     Ok(dsl::coalesce(&exprs).into())
@@ -103,10 +107,6 @@ pub fn collect_all(lfs: RArray) -> RbResult<RArray> {
     })))
 }
 
-pub fn cols(names: Vec<String>) -> RbExpr {
-    dsl::cols(names).into()
-}
-
 pub fn concat_lf(
     lfs: Value,
     rechunk: bool,
@@ -166,20 +166,24 @@ pub fn cum_fold(
     acc: &RbExpr,
     lambda: Value,
     exprs: RArray,
+    returns_scalar: bool,
+    return_dtype: Option<&RbDataTypeExpr>,
     include_init: bool,
 ) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
     let lambda = Opaque::from(lambda);
-
-    let func = move |a: Column, b: Column| {
-        binary_lambda(
-            Ruby::get().unwrap().get_inner(lambda),
-            a.take_materialized_series(),
-            b.take_materialized_series(),
-        )
-        .map(|v| v.map(Column::from))
-    };
-    Ok(dsl::cum_fold_exprs(acc.inner.clone(), func, exprs, include_init).into())
+    let func = PlanCallback::new(move |(a, b): (Series, Series)| {
+        binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b).map(|v| v.unwrap())
+    });
+    Ok(dsl::cum_fold_exprs(
+        acc.inner.clone(),
+        func,
+        exprs,
+        returns_scalar,
+        return_dtype.map(|v| v.inner.clone()),
+        include_init,
+    )
+    .into())
 }
 
 pub fn concat_lf_diagonal(
@@ -220,24 +224,6 @@ pub fn concat_lf_horizontal(lfs: RArray, parallel: bool) -> RbResult<RbLazyFrame
     Ok(lf.into())
 }
 
-pub fn dtype_cols(dtypes: RArray) -> RbResult<RbExpr> {
-    let dtypes = dtypes
-        .into_iter()
-        .map(Wrap::<DataType>::try_convert)
-        .collect::<RbResult<Vec<Wrap<DataType>>>>()?;
-    let dtypes = vec_extract_wrapped(dtypes);
-    Ok(dsl::dtype_cols(dtypes).into())
-}
-
-pub fn index_cols(indices: Vec<i64>) -> RbExpr {
-    if indices.len() == 1 {
-        dsl::nth(indices[0])
-    } else {
-        dsl::index_cols(indices)
-    }
-    .into()
-}
-
 #[allow(clippy::too_many_arguments)]
 pub fn duration(
     weeks: Option<&RbExpr>,
@@ -274,42 +260,28 @@ pub fn duration(
     dsl::duration(args).into()
 }
 
-pub fn first() -> RbExpr {
-    dsl::first().into()
-}
-
 pub fn fold(
     acc: &RbExpr,
     lambda: Value,
     exprs: RArray,
     returns_scalar: bool,
-    return_dtype: Option<Wrap<DataType>>,
+    return_dtype: Option<&RbDataTypeExpr>,
 ) -> RbResult<RbExpr> {
     let exprs = rb_exprs_to_exprs(exprs)?;
     let lambda = Opaque::from(lambda);
-
-    let func = move |a: Column, b: Column| {
-        binary_lambda(
-            Ruby::get().unwrap().get_inner(lambda),
-            a.take_materialized_series(),
-            b.take_materialized_series(),
-        )
-        .map(|v| v.map(Column::from))
-    };
+    let func = PlanCallback::new(move |(a, b): (Series, Series)| {
+        binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b).map(|v| v.unwrap())
+    });
     Ok(dsl::fold_exprs(
         acc.inner.clone(),
         func,
         exprs,
         returns_scalar,
-        return_dtype.map(|w| w.0),
+        return_dtype.map(|w| w.inner.clone()),
     )
     .into())
 }
 
-pub fn last() -> RbExpr {
-    dsl::last().into()
-}
-
 pub fn lit(value: Value, allow_object: bool, is_scalar: bool) -> RbResult<RbExpr> {
     if value.is_kind_of(class::true_class()) || value.is_kind_of(class::false_class()) {
         Ok(dsl::lit(bool::try_convert(value)?).into())

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

@@ -1,8 +1,8 @@
 use magnus::{IntoValue, Value};
 use polars_core;
+use polars_core::POOL;
 use polars_core::fmt::FloatFmt;
 use polars_core::prelude::IDX_DTYPE;
-use polars_core::POOL;
 
 use crate::conversion::Wrap;
 use crate::{RbResult, RbValueError};
@@ -22,7 +22,7 @@ pub fn set_float_fmt(fmt: String) -> RbResult<()> {
         e => {
             return Err(RbValueError::new_err(format!(
                 "fmt must be one of {{'full', 'mixed'}}, got {e}",
-            )))
+            )));
         }
     };
     polars_core::fmt::set_float_fmt(fmt);

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

@@ -1,6 +1,6 @@
+use crate::RbResult;
 use crate::conversion::Wrap;
 use crate::prelude::DataType;
-use crate::RbResult;
 
 pub fn dtype_str_repr(dtype: Wrap<DataType>) -> RbResult<String> {
     let dtype = dtype.0;

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

@@ -1,17 +1,17 @@
 use crate::RbResult;
 use magnus::{RArray, Ruby, Value};
-use polars_core::StringCacheHolder;
 
 pub fn enable_string_cache() {
-    polars_core::enable_string_cache()
+    // The string cache no longer exists.
 }
 
 pub fn disable_string_cache() {
-    polars_core::disable_string_cache()
+    // The string cache no longer exists.
 }
 
 pub fn using_string_cache() -> bool {
-    polars_core::using_string_cache()
+    // The string cache no longer exists.
+    true
 }
 
 #[magnus::wrap(class = "Polars::RbStringCacheHolder")]
@@ -19,7 +19,6 @@ pub struct RbStringCacheHolder {}
 
 impl RbStringCacheHolder {
     pub fn hold() -> RbResult<Value> {
-        let _hold = StringCacheHolder::hold();
         Ruby::get().unwrap().yield_splat(RArray::new())
     }
 }

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

@@ -1,4 +1,4 @@
-use magnus::{class, prelude::*, IntoValue, Module, RArray, RClass, RModule, Value};
+use magnus::{IntoValue, Module, RArray, RClass, RModule, Value, class, prelude::*};
 
 use crate::RbResult;
 

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

@@ -3,10 +3,10 @@ use num_traits::{Float, NumCast};
 use polars_core::prelude::*;
 
 use super::numo_rs::{Element, RbArray1};
+use crate::RbResult;
 use crate::error::RbPolarsErr;
 use crate::raise_err;
 use crate::series::RbSeries;
-use crate::RbResult;
 
 impl RbSeries {
     /// Convert this Series to a Numo array.

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

@@ -0,0 +1,102 @@
+use std::sync::Arc;
+
+use magnus::{TryConvert, Value, value::ReprValue};
+use polars::prelude::deletion::DeletionFilesList;
+use polars::prelude::{
+    CastColumnsPolicy, ColumnMapping, ExtraColumnsPolicy, MissingColumnsPolicy, PlSmallStr, Schema,
+    UnifiedScanArgs,
+};
+use polars_io::{HiveOptions, RowIndex};
+use polars_utils::IdxSize;
+use polars_utils::plpath::PlPathRef;
+use polars_utils::slice_enum::Slice;
+
+use crate::RbResult;
+use crate::prelude::Wrap;
+
+/// Interface to `class ScanOptions` on the Ruby side
+pub struct RbScanOptions(Value);
+
+impl TryConvert for RbScanOptions {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        Ok(Self(ob))
+    }
+}
+
+impl RbScanOptions {
+    pub fn extract_unified_scan_args(
+        &self,
+        // For cloud_options init
+        first_path: Option<PlPathRef>,
+    ) -> RbResult<UnifiedScanArgs> {
+        let row_index: Option<(Wrap<PlSmallStr>, IdxSize)> = self.0.funcall("row_index", ())?;
+        let pre_slice: Option<(i64, usize)> = self.0.funcall("pre_slice", ())?;
+        let cast_options: Wrap<CastColumnsPolicy> = self.0.funcall("cast_options", ())?;
+        let extra_columns: Wrap<ExtraColumnsPolicy> = self.0.funcall("extra_columns", ())?;
+        let missing_columns: Wrap<MissingColumnsPolicy> = self.0.funcall("missing_columns", ())?;
+        let include_file_paths: Option<Wrap<PlSmallStr>> =
+            self.0.funcall("include_file_paths", ())?;
+        let glob: bool = self.0.funcall("glob", ())?;
+        let hive_partitioning: Option<bool> = self.0.funcall("hive_partitioning", ())?;
+        let hive_schema: Option<Wrap<Schema>> = self.0.funcall("hive_schema", ())?;
+        let try_parse_hive_dates: bool = self.0.funcall("try_parse_hive_dates", ())?;
+        let rechunk: bool = self.0.funcall("rechunk", ())?;
+        let cache: bool = self.0.funcall("cache", ())?;
+        let storage_options: Option<Vec<(String, String)>> =
+            self.0.funcall("storage_options", ())?;
+        let retries: usize = self.0.funcall("retries", ())?;
+        let deletion_files: Option<Wrap<DeletionFilesList>> =
+            self.0.funcall("deletion_files", ())?;
+        let column_mapping: Option<Wrap<ColumnMapping>> = self.0.funcall("column_mapping", ())?;
+
+        let cloud_options = storage_options;
+
+        let cloud_options = if let Some(first_path) = first_path {
+            use crate::prelude::parse_cloud_options;
+
+            let first_path_url = first_path.to_str();
+            let cloud_options =
+                parse_cloud_options(first_path_url, cloud_options.unwrap_or_default())?;
+
+            Some(cloud_options.with_max_retries(retries))
+        } else {
+            None
+        };
+
+        let hive_schema = hive_schema.map(|s| Arc::new(s.0));
+
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: name.0,
+            offset,
+        });
+
+        let hive_options = HiveOptions {
+            enabled: hive_partitioning,
+            hive_start_idx: 0,
+            schema: hive_schema,
+            try_parse_dates: try_parse_hive_dates,
+        };
+
+        let unified_scan_args = UnifiedScanArgs {
+            // Schema is currently still stored inside the options per scan type, but we do eventually
+            // want to put it here instead.
+            schema: None,
+            cloud_options,
+            hive_options,
+            rechunk,
+            cache,
+            glob,
+            projection: None,
+            row_index,
+            pre_slice: pre_slice.map(Slice::from),
+            cast_columns_policy: cast_options.0,
+            missing_columns_policy: missing_columns.0,
+            extra_columns_policy: extra_columns.0,
+            include_file_paths: include_file_paths.map(|x| x.0),
+            deletion_files: DeletionFilesList::filter_empty(deletion_files.map(|x| x.0)),
+            column_mapping: column_mapping.map(|x| x.0),
+        };
+
+        Ok(unified_scan_args)
+    }
+}