polars-df 0.23.0 → 0.24.0

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

@@ -2,12 +2,15 @@ use std::io::BufReader;
 
 use magnus::{RHash, Ruby, Value};
 use polars::prelude::ArrowSchema;
+use polars::prelude::PlPathRef;
+use polars_io::cloud::CloudOptions;
 
 use crate::conversion::Wrap;
 use crate::file::{EitherRustRubyFile, get_either_file};
+use crate::utils::EnterPolarsExt;
 use crate::{RbPolarsErr, RbResult};
 
-pub fn read_ipc_schema(rb_f: Value) -> RbResult<RHash> {
+pub fn read_ipc_schema(rb: &Ruby, rb_f: Value) -> RbResult<RHash> {
     use arrow::io::ipc::read::read_file_metadata;
     let metadata = match get_either_file(rb_f, false)? {
         EitherRustRubyFile::Rust(r) => {
@@ -16,26 +19,70 @@ pub fn read_ipc_schema(rb_f: Value) -> RbResult<RHash> {
         EitherRustRubyFile::Rb(mut r) => read_file_metadata(&mut r).map_err(RbPolarsErr::from)?,
     };
 
-    let ruby = Ruby::get_with(rb_f);
-    let dict = ruby.hash_new();
+    let dict = rb.hash_new();
     fields_to_rbdict(&metadata.schema, &dict)?;
     Ok(dict)
 }
 
-pub fn read_parquet_metadata(rb_f: Value) -> RbResult<RHash> {
+pub fn read_parquet_metadata(
+    rb: &Ruby,
+    rb_f: Value,
+    storage_options: Option<Vec<(String, String)>>,
+    credential_provider: Option<Value>,
+    retries: usize,
+) -> RbResult<RHash> {
+    use std::io::Cursor;
+
+    use polars_io::pl_async::get_runtime;
     use polars_parquet::read::read_metadata;
     use polars_parquet::read::schema::read_custom_key_value_metadata;
+    use polars_utils::plpath::PlPath;
 
-    let metadata = match get_either_file(rb_f, false)? {
-        EitherRustRubyFile::Rust(r) => {
-            read_metadata(&mut BufReader::new(r)).map_err(RbPolarsErr::from)?
+    use crate::file::{RubyScanSourceInput, get_ruby_scan_source_input};
+
+    let metadata = match get_ruby_scan_source_input(rb_f, false)? {
+        RubyScanSourceInput::Buffer(buf) => {
+            read_metadata(&mut Cursor::new(buf)).map_err(RbPolarsErr::from)?
+        }
+        RubyScanSourceInput::Path(p) => {
+            let cloud_options = parse_cloud_options(
+                Some(p.as_ref()),
+                storage_options,
+                credential_provider,
+                retries,
+            )?;
+            match p {
+                PlPath::Local(local) => {
+                    let file = polars_utils::open_file(&local).map_err(RbPolarsErr::from)?;
+                    read_metadata(&mut BufReader::new(file)).map_err(RbPolarsErr::from)?
+                }
+                PlPath::Cloud(_) => {
+                    use polars::prelude::ParquetObjectStore;
+                    use polars_error::PolarsResult;
+
+                    rb.detach(|| {
+                        get_runtime().block_on(async {
+                            let mut reader = ParquetObjectStore::from_uri(
+                                p.as_ref(),
+                                cloud_options.as_ref(),
+                                None,
+                            )
+                            .await?;
+                            let result = reader.get_metadata().await?;
+                            PolarsResult::Ok((**result).clone())
+                        })
+                    })
+                    .map_err(RbPolarsErr::from)?
+                }
+            }
+        }
+        RubyScanSourceInput::File(f) => {
+            read_metadata(&mut BufReader::new(f)).map_err(RbPolarsErr::from)?
         }
-        EitherRustRubyFile::Rb(mut r) => read_metadata(&mut r).map_err(RbPolarsErr::from)?,
     };
 
     let key_value_metadata = read_custom_key_value_metadata(metadata.key_value_metadata());
-    let ruby = Ruby::get_with(rb_f);
-    let dict = ruby.hash_new();
+    let dict = rb.hash_new();
     for (key, value) in key_value_metadata.into_iter() {
         dict.aset(key.as_str(), value.as_str())?;
     }
@@ -66,3 +113,27 @@ fn fields_to_rbdict(schema: &ArrowSchema, dict: &RHash) -> RbResult<()> {
     }
     Ok(())
 }
+
+pub fn parse_cloud_options<'a>(
+    first_path: Option<PlPathRef<'a>>,
+    storage_options: Option<Vec<(String, String)>>,
+    credential_provider: Option<Value>,
+    retries: usize,
+) -> RbResult<Option<CloudOptions>> {
+    let result = 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, storage_options.unwrap_or_default())?;
+
+        Some(
+            cloud_options
+                .with_max_retries(retries)
+                .with_credential_provider(credential_provider.map(|_| todo!())),
+        )
+    } else {
+        None
+    };
+    Ok(result)
+}

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

@@ -6,9 +6,11 @@ use polars::lazy::dsl;
 use polars::prelude::*;
 
 use crate::conversion::{Wrap, get_lf, get_rbseq};
+use crate::expr::ToExprs;
 use crate::expr::datatype::RbDataTypeExpr;
+use crate::lazyframe::RbOptFlags;
 use crate::map::lazy::binary_lambda;
-use crate::rb_exprs_to_exprs;
+use crate::utils::{EnterPolarsExt, RubyAttach};
 use crate::{RbDataFrame, RbExpr, RbLazyFrame, RbPolarsErr, RbResult, RbSeries, RbValueError};
 
 macro_rules! set_unwrapped_or_0 {
@@ -62,7 +64,7 @@ pub fn arg_sort_by(
     multithreaded: bool,
     maintain_order: bool,
 ) -> RbResult<RbExpr> {
-    let by = rb_exprs_to_exprs(by)?;
+    let by = by.to_exprs()?;
     Ok(dsl::arg_sort_by(
         by,
         SortMultipleOptions {
@@ -81,7 +83,12 @@ pub fn arg_where(condition: &RbExpr) -> RbExpr {
 }
 
 pub fn as_struct(exprs: RArray) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
+    if exprs.is_empty() {
+        return Err(RbValueError::new_err(
+            "expected at least 1 expression in 'as_struct'",
+        ));
+    }
     Ok(dsl::as_struct(exprs).into())
 }
 
@@ -90,7 +97,7 @@ pub fn field(names: Vec<String>) -> RbExpr {
 }
 
 pub fn coalesce(exprs: RArray) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
     Ok(dsl::coalesce(&exprs).into())
 }
 
@@ -98,13 +105,25 @@ pub fn col(name: String) -> RbExpr {
     dsl::col(&name).into()
 }
 
-pub fn collect_all(ruby: &Ruby, lfs: RArray) -> RbResult<RArray> {
+fn lfs_to_plans(lfs: RArray) -> RbResult<Vec<DslPlan>> {
     let lfs = lfs.typecheck::<Obj<RbLazyFrame>>()?;
+    Ok(lfs
+        .into_iter()
+        .map(|lf| lf.ldf.read().logical_plan.clone())
+        .collect())
+}
 
-    Ok(ruby.ary_from_iter(lfs.iter().map(|lf| {
-        let df = lf.ldf.borrow().clone().collect().unwrap();
-        RbDataFrame::new(df)
-    })))
+pub fn collect_all(
+    ruby: &Ruby,
+    lfs: RArray,
+    engine: Wrap<Engine>,
+    optflags: &RbOptFlags,
+) -> RbResult<RArray> {
+    let plans = lfs_to_plans(lfs)?;
+    let dfs = ruby.enter_polars(|| {
+        LazyFrame::collect_all_with_engine(plans, engine.0, optflags.clone().inner.into_inner())
+    })?;
+    Ok(ruby.ary_from_iter(dfs.into_iter().map(Into::<RbDataFrame>::into)))
 }
 
 pub fn concat_lf(
@@ -136,13 +155,19 @@ pub fn concat_lf(
 }
 
 pub fn concat_list(s: RArray) -> RbResult<RbExpr> {
-    let s = rb_exprs_to_exprs(s)?;
+    let s = s.to_exprs()?;
     let expr = dsl::concat_list(s).map_err(RbPolarsErr::from)?;
     Ok(expr.into())
 }
 
+pub fn concat_arr(s: RArray) -> RbResult<RbExpr> {
+    let s = s.to_exprs()?;
+    let expr = dsl::concat_arr(s).map_err(RbPolarsErr::from)?;
+    Ok(expr.into())
+}
+
 pub fn concat_str(s: RArray, separator: String, ignore_nulls: bool) -> RbResult<RbExpr> {
-    let s = rb_exprs_to_exprs(s)?;
+    let s = s.to_exprs()?;
     Ok(dsl::concat_str(s, &separator, ignore_nulls).into())
 }
 
@@ -158,10 +183,6 @@ pub fn arctan2(y: &RbExpr, x: &RbExpr) -> RbExpr {
     y.inner.clone().arctan2(x.inner.clone()).into()
 }
 
-pub fn arctan2d(y: &RbExpr, x: &RbExpr) -> RbExpr {
-    y.inner.clone().arctan2(x.inner.clone()).degrees().into()
-}
-
 pub fn cum_fold(
     acc: &RbExpr,
     lambda: Value,
@@ -170,10 +191,10 @@ pub fn cum_fold(
     return_dtype: Option<&RbDataTypeExpr>,
     include_init: bool,
 ) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
     let lambda = Opaque::from(lambda);
     let func = PlanCallback::new(move |(a, b): (Series, Series)| {
-        binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b).map(|v| v.unwrap())
+        Ruby::attach(|rb| binary_lambda(rb.get_inner(lambda), a, b).map(|v| v.unwrap()))
     });
     Ok(dsl::cum_fold_exprs(
         acc.inner.clone(),
@@ -186,6 +207,60 @@ pub fn cum_fold(
     .into())
 }
 
+pub fn cum_reduce(
+    lambda: Value,
+    exprs: RArray,
+    returns_scalar: bool,
+    return_dtype: Option<&RbDataTypeExpr>,
+) -> RbResult<RbExpr> {
+    let exprs = exprs.to_exprs()?;
+    let lambda = Opaque::from(lambda);
+    let func = PlanCallback::new(move |(a, b): (Series, Series)| {
+        Ruby::attach(|rb| binary_lambda(rb.get_inner(lambda), a, b).map(|v| v.unwrap()))
+    });
+    Ok(dsl::cum_reduce_exprs(
+        func,
+        exprs,
+        returns_scalar,
+        return_dtype.map(|v| v.inner.clone()),
+    )
+    .into())
+}
+
+pub fn datetime(
+    year: &RbExpr,
+    month: &RbExpr,
+    day: &RbExpr,
+    hour: Option<&RbExpr>,
+    minute: Option<&RbExpr>,
+    second: Option<&RbExpr>,
+    microsecond: Option<&RbExpr>,
+    time_unit: Wrap<TimeUnit>,
+    time_zone: Wrap<Option<TimeZone>>,
+    ambiguous: &RbExpr,
+) -> RbExpr {
+    let year = year.inner.clone();
+    let month = month.inner.clone();
+    let day = day.inner.clone();
+    set_unwrapped_or_0!(hour, minute, second, microsecond);
+    let ambiguous = ambiguous.inner.clone();
+    let time_unit = time_unit.0;
+    let time_zone = time_zone.0;
+    let args = DatetimeArgs {
+        year,
+        month,
+        day,
+        hour,
+        minute,
+        second,
+        microsecond,
+        time_unit,
+        time_zone,
+        ambiguous,
+    };
+    dsl::datetime(args).into()
+}
+
 pub fn concat_lf_diagonal(
     lfs: RArray,
     rechunk: bool,
@@ -224,7 +299,6 @@ pub fn concat_lf_horizontal(lfs: RArray, parallel: bool) -> RbResult<RbLazyFrame
     Ok(lf.into())
 }
 
-#[allow(clippy::too_many_arguments)]
 pub fn duration(
     weeks: Option<&RbExpr>,
     days: Option<&RbExpr>,
@@ -267,10 +341,10 @@ pub fn fold(
     returns_scalar: bool,
     return_dtype: Option<&RbDataTypeExpr>,
 ) -> RbResult<RbExpr> {
-    let exprs = rb_exprs_to_exprs(exprs)?;
+    let exprs = exprs.to_exprs()?;
     let lambda = Opaque::from(lambda);
     let func = PlanCallback::new(move |(a, b): (Series, Series)| {
-        binary_lambda(Ruby::get().unwrap().get_inner(lambda), a, b).map(|v| v.unwrap())
+        Ruby::attach(|rb| binary_lambda(rb.get_inner(lambda), a, b).map(|v| v.unwrap()))
     });
     Ok(dsl::fold_exprs(
         acc.inner.clone(),
@@ -309,7 +383,7 @@ pub fn lit(value: Value, allow_object: bool, is_scalar: bool) -> RbResult<RbExpr
             Ok(dsl::lit(unsafe { v.as_slice() }).into())
         }
     } else if let Ok(series) = Obj::<RbSeries>::try_convert(value) {
-        let s = series.series.borrow();
+        let s = series.series.read();
         if is_scalar {
             let av = s
                 .get(0)

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

@@ -8,4 +8,6 @@ pub mod misc;
 pub mod random;
 pub mod range;
 pub mod string_cache;
+pub mod strings;
+pub mod utils;
 pub mod whenthen;

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, RbResult};
+use crate::{RbDataTypeExpr, RbExpr, RbPolarsErr, RbResult};
 
 pub fn int_range(start: &RbExpr, end: &RbExpr, step: i64, dtype: Wrap<DataType>) -> RbExpr {
     let start = start.inner.clone();
@@ -16,9 +16,9 @@ pub fn int_ranges(
     start: &RbExpr,
     end: &RbExpr,
     step: &RbExpr,
-    dtype: Wrap<DataType>,
+    dtype: &RbDataTypeExpr,
 ) -> RbResult<RbExpr> {
-    let dtype = dtype.0;
+    let dtype = dtype.inner.clone();
     Ok(dsl::int_ranges(
         start.inner.clone(),
         end.inner.clone(),
@@ -113,3 +113,19 @@ pub fn time_ranges(
     let closed = closed.0;
     dsl::time_ranges(start, end, every, closed).into()
 }
+
+pub fn linear_spaces(
+    start: &RbExpr,
+    end: &RbExpr,
+    num_samples: &RbExpr,
+    closed: Wrap<ClosedInterval>,
+    as_array: bool,
+) -> RbResult<RbExpr> {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let num_samples = num_samples.inner.clone();
+    let closed = closed.0;
+    let out =
+        dsl::linear_spaces(start, end, num_samples, closed, as_array).map_err(RbPolarsErr::from)?;
+    Ok(out.into())
+}

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

@@ -0,0 +1,6 @@
+use crate::RbResult;
+
+pub fn escape_regex(s: String) -> RbResult<String> {
+    let escaped_s = polars_ops::chunked_array::strings::escape_regex_str(&s);
+    Ok(escaped_s)
+}

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

@@ -0,0 +1,6 @@
+use crate::RbResult;
+use polars_core::config::get_engine_affinity;
+
+pub fn rb_get_engine_affinity() -> RbResult<String> {
+    Ok(get_engine_affinity())
+}

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

@@ -1 +1,50 @@
-pub mod to_ruby;
+use magnus::prelude::*;
+use magnus::{RHash, Value};
+use polars::prelude::{ArrowDataType, DataType};
+use polars_error::polars_err;
+
+use crate::interop::arrow::to_rust::normalize_arrow_fields;
+use crate::prelude::Wrap;
+use crate::series::import_schema_rbcapsule;
+use crate::utils::to_rb_err;
+use crate::{RbResult, RbValueError};
+
+pub mod to_rb;
+pub mod to_rust;
+
+pub fn init_polars_schema_from_arrow_c_schema(
+    polars_schema: RHash,
+    schema_object: Value,
+) -> RbResult<()> {
+    let schema_capsule = schema_object.funcall("arrow_c_schema", ())?;
+
+    let field = import_schema_rbcapsule(schema_capsule)?;
+    let field = normalize_arrow_fields(&field);
+
+    let ArrowDataType::Struct(fields) = field.dtype else {
+        return Err(RbValueError::new_err(format!(
+            "arrow_c_schema of object passed to Polars::Schema did not return struct dtype: \
+            object: {}, dtype: {:?}",
+            schema_object, &field.dtype
+        )));
+    };
+
+    for field in fields {
+        let dtype = DataType::from_arrow_field(&field);
+
+        let name = field.name.as_str();
+        let dtype = Wrap(dtype);
+
+        if polars_schema.get(name).is_some() {
+            return Err(to_rb_err(polars_err!(
+                Duplicate:
+                "arrow schema contained duplicate name: {}",
+                name
+            )));
+        }
+
+        polars_schema.aset(name, dtype)?;
+    }
+
+    Ok(())
+}

data/ext/polars/src/interop/arrow/{to_ruby.rs → to_rb.rs}

@@ -20,6 +20,17 @@ impl RbArrowArrayStream {
     }
 }
 
+#[magnus::wrap(class = "Polars::ArrowSchema")]
+pub struct RbArrowSchema {
+    pub(crate) schema: ffi::ArrowSchema,
+}
+
+impl RbArrowSchema {
+    pub fn to_i(&self) -> usize {
+        (&self.schema as *const _) as usize
+    }
+}
+
 pub(crate) fn dataframe_to_stream(df: &DataFrame, ruby: &Ruby) -> RbResult<Value> {
     let iter = Box::new(DataFrameStreamIterator::new(df));
     let field = iter.field();
@@ -27,6 +38,25 @@ pub(crate) fn dataframe_to_stream(df: &DataFrame, ruby: &Ruby) -> RbResult<Value
     Ok(RbArrowArrayStream { stream }.into_value_with(ruby))
 }
 
+pub(crate) fn polars_schema_to_rbcapsule(
+    ruby: &Ruby,
+    schema: crate::prelude::Wrap<polars::prelude::Schema>,
+) -> RbResult<Value> {
+    let schema: arrow::ffi::ArrowSchema = arrow::ffi::export_field_to_c(&ArrowField::new(
+        PlSmallStr::EMPTY,
+        ArrowDataType::Struct(
+            schema
+                .0
+                .iter_fields()
+                .map(|x| x.to_arrow(CompatLevel::newest()))
+                .collect(),
+        ),
+        false,
+    ));
+
+    Ok(RbArrowSchema { schema }.into_value_with(ruby))
+}
+
 pub struct DataFrameStreamIterator {
     columns: Vec<Series>,
     dtype: ArrowDataType,

data/ext/polars/src/interop/arrow/to_rust.rs

@@ -0,0 +1,43 @@
+use polars_core::prelude::*;
+
+pub(crate) fn normalize_arrow_fields(field: &ArrowField) -> ArrowField {
+    // normalize fields with extension dtypes that are otherwise standard dtypes associated
+    // with (for us) irrelevant metadata; recreate the field using the inner (standard) dtype
+    match field {
+        ArrowField {
+            dtype: ArrowDataType::Struct(fields),
+            ..
+        } => {
+            let mut normalized = false;
+            let normalized_fields: Vec<_> = fields
+                .iter()
+                .map(|f| {
+                    // note: google bigquery column data is returned as a standard arrow dtype, but the
+                    // sql type it was loaded from is associated as metadata (resulting in an extension dtype)
+                    if let ArrowDataType::Extension(ext_type) = &f.dtype
+                        && ext_type.name.starts_with("google:sqlType:")
+                    {
+                        normalized = true;
+                        return ArrowField::new(
+                            f.name.clone(),
+                            ext_type.inner.clone(),
+                            f.is_nullable,
+                        );
+                    }
+                    f.clone()
+                })
+                .collect();
+
+            if normalized {
+                ArrowField::new(
+                    field.name.clone(),
+                    ArrowDataType::Struct(normalized_fields),
+                    field.is_nullable,
+                )
+            } else {
+                field.clone()
+            }
+        }
+        _ => field.clone(),
+    }
+}

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

@@ -6,7 +6,7 @@ use crate::dataframe::RbDataFrame;
 impl RbDataFrame {
     pub fn to_numo(&self) -> Option<Value> {
         let mut st = None;
-        for s in self.df.borrow().iter() {
+        for s in self.df.read().iter() {
             let dt_i = s.dtype();
             match st {
                 None => st = Some(dt_i.clone()),

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

@@ -11,7 +11,7 @@ use crate::series::RbSeries;
 impl RbSeries {
     /// Convert this Series to a Numo array.
     pub fn to_numo(&self) -> RbResult<Value> {
-        series_to_numo(&self.series.borrow())
+        series_to_numo(&self.series.read())
     }
 }
 

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

@@ -0,0 +1,39 @@
+use magnus::Ruby;
+use polars::prelude::*;
+
+use super::RbLazyFrame;
+use crate::utils::EnterPolarsExt;
+use crate::{RbDataFrame, RbResult};
+
+impl RbLazyFrame {
+    pub fn collect_concurrently(ruby: &Ruby, self_: &Self) -> RbResult<RbInProcessQuery> {
+        let ipq = ruby.enter_polars(|| {
+            let ldf = self_.ldf.read().clone();
+            ldf.collect_concurrently()
+        })?;
+        Ok(RbInProcessQuery { ipq })
+    }
+}
+
+#[magnus::wrap(class = "Polars::RbInProcessQuery")]
+#[repr(transparent)]
+#[derive(Clone)]
+pub struct RbInProcessQuery {
+    pub ipq: InProcessQuery,
+}
+
+impl RbInProcessQuery {
+    pub fn cancel(ruby: &Ruby, self_: &Self) -> RbResult<()> {
+        ruby.enter_polars_ok(|| self_.ipq.cancel())
+    }
+
+    pub fn fetch(ruby: &Ruby, self_: &Self) -> RbResult<Option<RbDataFrame>> {
+        let out = ruby.enter_polars(|| self_.ipq.fetch().transpose())?;
+        Ok(out.map(|df| df.into()))
+    }
+
+    pub fn fetch_blocking(ruby: &Ruby, self_: &Self) -> RbResult<RbDataFrame> {
+        let out = ruby.enter_polars(|| self_.ipq.fetch_blocking())?;
+        Ok(out.into())
+    }
+}