polars-df 0.19.0 → 0.21.0

data/ext/polars/src/file.rs

@@ -3,17 +3,18 @@ use std::io;
 use std::io::{Cursor, Read, Seek, SeekFrom, Write};
 use std::path::PathBuf;
 
-use magnus::{exception, prelude::*, value::Opaque, Error, RString, Ruby, Value};
+use magnus::{Error, RString, Ruby, Value, exception, prelude::*, value::Opaque};
 use polars::io::cloud::CloudOptions;
 use polars::io::mmap::MmapBytesReader;
+use polars::prelude::PlPath;
 use polars::prelude::file::DynWriteable;
 use polars::prelude::sync_on_close::SyncOnCloseType;
 use polars_utils::file::ClosableFile;
 use polars_utils::mmap::MemSlice;
 
+use crate::RbResult;
 use crate::error::RbPolarsErr;
 use crate::prelude::resolve_homedir;
-use crate::RbResult;
 
 #[derive(Clone)]
 pub struct RbFileLikeObject {
@@ -93,7 +94,7 @@ impl RbFileLikeObject {
 
 /// Extracts a string repr from, and returns an IO error to send back to rust.
 fn rberr_to_io_err(e: Error) -> io::Error {
-    io::Error::new(io::ErrorKind::Other, e.to_string())
+    io::Error::other(e.to_string())
 }
 
 impl Read for RbFileLikeObject {
@@ -188,7 +189,7 @@ impl EitherRustRubyFile {
 
 pub enum RubyScanSourceInput {
     Buffer(MemSlice),
-    Path(PathBuf),
+    Path(PlPath),
     #[allow(dead_code)]
     File(File),
 }
@@ -202,8 +203,13 @@ pub(crate) fn try_get_rbfile(
 }
 
 pub fn get_ruby_scan_source_input(rb_f: Value, write: bool) -> RbResult<RubyScanSourceInput> {
-    if let Ok(file_path) = PathBuf::try_convert(rb_f) {
-        // TODO resolve_homedir
+    if let Ok(s) = String::try_convert(rb_f) {
+        let mut file_path = PlPath::new(&s);
+        if let Some(p) = file_path.as_ref().as_local_path() {
+            if p.starts_with("~/") {
+                file_path = PlPath::Local(resolve_homedir(&p).into());
+            }
+        }
         Ok(RubyScanSourceInput::Path(file_path))
     } else {
         let f = RbFileLikeObject::with_requirements(rb_f, !write, write, !write)?;

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

@@ -1,13 +1,10 @@
 use std::io::BufReader;
 
-use arrow::array::Utf8ViewArray;
 use magnus::{RHash, Value};
 use polars::prelude::ArrowSchema;
-use polars_core::datatypes::create_enum_dtype;
 
 use crate::conversion::Wrap;
-use crate::file::{get_either_file, EitherRustRubyFile};
-use crate::prelude::ArrowDataType;
+use crate::file::{EitherRustRubyFile, get_either_file};
 use crate::{RbPolarsErr, RbResult};
 
 pub fn read_ipc_schema(rb_f: Value) -> RbResult<RHash> {
@@ -42,13 +39,7 @@ pub fn read_parquet_schema(rb_f: Value) -> RbResult<RHash> {
 
 fn fields_to_rbdict(schema: &ArrowSchema, dict: &RHash) -> RbResult<()> {
     for field in schema.iter_values() {
-        let dt = if field.is_enum() {
-            Wrap(create_enum_dtype(Utf8ViewArray::new_empty(
-                ArrowDataType::Utf8View,
-            )))
-        } else {
-            Wrap(polars::prelude::DataType::from_arrow_field(field))
-        };
+        let dt = Wrap(polars::prelude::DataType::from_arrow_field(field));
         dict.aset(field.name.as_str(), dt)?;
     }
     Ok(())

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};
 
@@ -103,10 +103,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 +162,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 +220,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 +256,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/range.rs

@@ -3,7 +3,7 @@ use polars_core::datatypes::{TimeUnit, TimeZone};
 
 use crate::conversion::Wrap;
 use crate::prelude::*;
-use crate::RbExpr;
+use crate::{RbExpr, RbResult};
 
 pub fn int_range(start: &RbExpr, end: &RbExpr, step: i64, dtype: Wrap<DataType>) -> RbExpr {
     let start = start.inner.clone();
@@ -12,16 +12,20 @@ pub fn int_range(start: &RbExpr, end: &RbExpr, step: i64, dtype: Wrap<DataType>)
     dsl::int_range(start, end, step, dtype).into()
 }
 
-pub fn int_ranges(start: &RbExpr, end: &RbExpr, step: &RbExpr, dtype: Wrap<DataType>) -> RbExpr {
+pub fn int_ranges(
+    start: &RbExpr,
+    end: &RbExpr,
+    step: &RbExpr,
+    dtype: Wrap<DataType>,
+) -> RbResult<RbExpr> {
     let dtype = dtype.0;
-
-    let mut result = dsl::int_ranges(start.inner.clone(), end.inner.clone(), step.inner.clone());
-
-    if dtype != DataType::Int64 {
-        result = result.cast(DataType::List(Box::new(dtype)))
-    }
-
-    result.into()
+    Ok(dsl::int_ranges(
+        start.inner.clone(),
+        end.inner.clone(),
+        step.inner.clone(),
+        dtype,
+    )
+    .into())
 }
 
 pub fn date_range(

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)
+    }
+}