polars-df 0.15.0 → 0.26.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::{RbDataTypeExpr, RbExpr, RbPolarsErr, 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 {
-    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()
+pub fn int_ranges(
+    start: &RbExpr,
+    end: &RbExpr,
+    step: &RbExpr,
+    dtype: &RbDataTypeExpr,
+) -> RbResult<RbExpr> {
+    let dtype = dtype.inner.clone();
+    Ok(dsl::int_ranges(
+        start.inner.clone(),
+        end.inner.clone(),
+        step.inner.clone(),
+        dtype,
+    )
+    .into())
 }
 
 pub fn date_range(
@@ -29,12 +33,14 @@ pub fn date_range(
     end: &RbExpr,
     interval: String,
     closed: Wrap<ClosedWindow>,
-) -> RbExpr {
+) -> RbResult<RbExpr> {
     let start = start.inner.clone();
     let end = end.inner.clone();
     let interval = Duration::parse(&interval);
     let closed = closed.0;
-    dsl::date_range(start, end, interval, closed).into()
+    let out = dsl::date_range(Some(start), Some(end), Some(interval), None, closed)
+        .map_err(RbPolarsErr::from)?;
+    Ok(out.into())
 }
 
 pub fn date_ranges(
@@ -42,46 +48,68 @@ pub fn date_ranges(
     end: &RbExpr,
     interval: String,
     closed: Wrap<ClosedWindow>,
-) -> RbExpr {
+) -> RbResult<RbExpr> {
     let start = start.inner.clone();
     let end = end.inner.clone();
     let interval = Duration::parse(&interval);
     let closed = closed.0;
-    dsl::date_ranges(start, end, interval, closed).into()
+    let out = dsl::date_ranges(Some(start), Some(end), Some(interval), None, closed)
+        .map_err(RbPolarsErr::from)?;
+    Ok(out.into())
 }
 
 pub fn datetime_range(
     start: &RbExpr,
     end: &RbExpr,
-    every: String,
+    interval: String,
     closed: Wrap<ClosedWindow>,
     time_unit: Option<Wrap<TimeUnit>>,
-    time_zone: Option<Wrap<TimeZone>>,
-) -> RbExpr {
+    time_zone: Wrap<Option<TimeZone>>,
+) -> RbResult<RbExpr> {
     let start = start.inner.clone();
     let end = end.inner.clone();
-    let every = Duration::parse(&every);
+    let interval = Duration::try_parse(&interval).map_err(RbPolarsErr::from)?;
     let closed = closed.0;
     let time_unit = time_unit.map(|x| x.0);
-    let time_zone = time_zone.map(|x| x.0);
-    dsl::datetime_range(start, end, every, closed, time_unit, time_zone).into()
+    let time_zone = time_zone.0;
+    let out = dsl::datetime_range(
+        Some(start),
+        Some(end),
+        Some(interval),
+        None,
+        closed,
+        time_unit,
+        time_zone,
+    )
+    .map_err(RbPolarsErr::from)?;
+    Ok(out.into())
 }
 
 pub fn datetime_ranges(
     start: &RbExpr,
     end: &RbExpr,
-    every: String,
+    interval: String,
     closed: Wrap<ClosedWindow>,
     time_unit: Option<Wrap<TimeUnit>>,
-    time_zone: Option<Wrap<TimeZone>>,
-) -> RbExpr {
+    time_zone: Wrap<Option<TimeZone>>,
+) -> RbResult<RbExpr> {
     let start = start.inner.clone();
     let end = end.inner.clone();
-    let every = Duration::parse(&every);
+    let interval = Duration::try_parse(&interval).map_err(RbPolarsErr::from)?;
     let closed = closed.0;
     let time_unit = time_unit.map(|x| x.0);
-    let time_zone = time_zone.map(|x| x.0);
-    dsl::datetime_ranges(start, end, every, closed, time_unit, time_zone).into()
+    let time_zone = time_zone.0;
+    let out = dsl::datetime_ranges(
+        Some(start),
+        Some(end),
+        Some(interval),
+        None,
+        closed,
+        time_unit,
+        time_zone,
+    )
+    .map_err(RbPolarsErr::from)?;
+    Ok(out.into())
 }
 
 pub fn time_range(
@@ -89,12 +117,12 @@ pub fn time_range(
     end: &RbExpr,
     every: String,
     closed: Wrap<ClosedWindow>,
-) -> RbExpr {
+) -> RbResult<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()
+    Ok(dsl::time_range(start, end, every, closed).into())
 }
 
 pub fn time_ranges(
@@ -102,10 +130,39 @@ pub fn time_ranges(
     end: &RbExpr,
     every: String,
     closed: Wrap<ClosedWindow>,
-) -> RbExpr {
+) -> RbResult<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()
+    Ok(dsl::time_ranges(start, end, every, closed).into())
+}
+
+pub fn linear_space(
+    start: &RbExpr,
+    end: &RbExpr,
+    num_samples: &RbExpr,
+    closed: Wrap<ClosedInterval>,
+) -> RbResult<RbExpr> {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let num_samples = num_samples.inner.clone();
+    let closed = closed.0;
+    Ok(dsl::linear_space(start, end, num_samples, 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,8 @@
+use crate::RbResult;
+
+pub fn rb_get_engine_affinity() -> RbResult<String> {
+    Ok(polars_config::config()
+        .engine_affinity()
+        .as_static_str()
+        .to_string())
+}

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

@@ -1 +1,52 @@
-pub mod to_ruby;
+use magnus::prelude::*;
+use magnus::{RHash, Ruby, 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::ruby::utils::TryIntoValue;
+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 ruby = &Ruby::get_with(polars_schema);
+    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).try_into_value_with(ruby)?;
+
+        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}

@@ -1,6 +1,6 @@
 use arrow::datatypes::ArrowDataType;
 use arrow::ffi;
-use magnus::{IntoValue, Value};
+use magnus::{IntoValue, Ruby, Value};
 use polars::datatypes::CompatLevel;
 use polars::frame::DataFrame;
 use polars::prelude::{ArrayRef, ArrowField, PlSmallStr, PolarsResult, SchemaExt};
@@ -9,9 +9,9 @@ use polars_core::utils::arrow;
 
 use crate::RbResult;
 
-#[magnus::wrap(class = "Polars::RbArrowArrayStream")]
+#[magnus::wrap(class = "Polars::ArrowArrayStream")]
 pub struct RbArrowArrayStream {
-    stream: ffi::ArrowArrayStream,
+    pub(crate) stream: ffi::ArrowArrayStream,
 }
 
 impl RbArrowArrayStream {
@@ -20,11 +20,41 @@ impl RbArrowArrayStream {
     }
 }
 
-pub(crate) fn dataframe_to_stream(df: &DataFrame) -> RbResult<Value> {
+#[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();
     let stream = ffi::export_iterator(iter, field);
-    Ok(RbArrowArrayStream { stream }.into_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 {
@@ -41,13 +71,13 @@ impl DataFrameStreamIterator {
 
         Self {
             columns: df
-                .get_columns()
+                .columns()
                 .iter()
                 .map(|v| v.as_materialized_series().clone())
                 .collect(),
             dtype,
             idx: 0,
-            n_chunks: df.n_chunks(),
+            n_chunks: df.first_col_n_chunks(),
         }
     }
 

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().columns().iter() {
             let dt_i = s.dtype();
             match st {
                 None => st = Some(dt_i.clone()),

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

@@ -1,61 +1,83 @@
-use magnus::{class, prelude::*, Module, RArray, RClass, RModule, Value};
-use polars::series::BitRepr;
+use magnus::{Ruby, Value};
+use num_traits::{Float, NumCast};
 use polars_core::prelude::*;
 
+use crate::RbResult;
 use crate::error::RbPolarsErr;
 use crate::raise_err;
+use crate::ruby::numo::{Element, RbArray1};
 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();
+    /// Convert this Series to a Numo array.
+    pub fn to_numo(rb: &Ruby, self_: &Self) -> RbResult<Value> {
+        series_to_numo(rb, &self_.series.read())
+    }
+}
 
-                // 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 let Some(BitRepr::Large(_)) = s.bit_repr() {
-                    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
-                );
-            }
+/// Convert a Series to a Numo array.
+fn series_to_numo(rb: &Ruby, s: &Series) -> RbResult<Value> {
+    series_to_numo_with_copy(rb, s)
+}
+
+/// Convert a Series to a Numo array, copying data in the process.
+fn series_to_numo_with_copy(rb: &Ruby, s: &Series) -> RbResult<Value> {
+    use DataType::*;
+    match s.dtype() {
+        Int8 => numeric_series_to_numo::<Int8Type, f32>(rb, s),
+        Int16 => numeric_series_to_numo::<Int16Type, f32>(rb, s),
+        Int32 => numeric_series_to_numo::<Int32Type, f64>(rb, s),
+        Int64 => numeric_series_to_numo::<Int64Type, f64>(rb, s),
+        UInt8 => numeric_series_to_numo::<UInt8Type, f32>(rb, s),
+        UInt16 => numeric_series_to_numo::<UInt16Type, f32>(rb, s),
+        UInt32 => numeric_series_to_numo::<UInt32Type, f64>(rb, s),
+        UInt64 => numeric_series_to_numo::<UInt64Type, f64>(rb, s),
+        Float32 => numeric_series_to_numo::<Float32Type, f32>(rb, s),
+        Float64 => numeric_series_to_numo::<Float64Type, f64>(rb, s),
+        Boolean => boolean_series_to_numo(rb, s),
+        String => {
+            let ca = s.str().unwrap();
+            let values = ca.iter();
+            RbArray1::from_iter(rb, values)
+        }
+        dt => {
+            raise_err!(
+                format!("'to_numo' not supported for dtype: {dt:?}"),
+                ComputeError
+            );
         }
     }
 }
+
+/// Convert numeric types to f32 or f64 with NaN representing a null value.
+fn numeric_series_to_numo<T, U>(rb: &Ruby, s: &Series) -> RbResult<Value>
+where
+    T: PolarsNumericType,
+    T::Native: Element,
+    U: Float + Element,
+{
+    let ca: &ChunkedArray<T> = s.as_ref().as_ref();
+    if s.null_count() == 0 {
+        let values = ca.into_no_null_iter();
+        RbArray1::<T::Native>::from_iter(rb, values)
+    } else {
+        let mapper = |opt_v: Option<T::Native>| match opt_v {
+            Some(v) => NumCast::from(v).unwrap(),
+            None => U::nan(),
+        };
+        let values = ca.iter().map(mapper);
+        RbArray1::from_iter(rb, values)
+    }
+}
+
+/// Convert booleans to bit if no nulls are present, otherwise convert to objects.
+fn boolean_series_to_numo(rb: &Ruby, s: &Series) -> RbResult<Value> {
+    let ca = s.bool().unwrap();
+    if s.null_count() == 0 {
+        let values = ca.no_null_iter();
+        RbArray1::<bool>::from_iter(rb, values)
+    } else {
+        let values = ca.iter();
+        RbArray1::from_iter(rb, values)
+    }
+}

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

@@ -0,0 +1,107 @@
+use std::time::Duration;
+
+use magnus::{RHash, TryConvert, Value, r_hash::ForEach, value::ReprValue};
+use polars::prelude::CloudScheme;
+use polars_io::cloud::{CloudOptions, CloudRetryConfig};
+use polars_utils::total_ord::TotalOrdWrap;
+
+use crate::utils::to_rb_err;
+use crate::{RbResult, RbValueError};
+
+pub struct OptRbCloudOptions(Value);
+
+impl TryConvert for OptRbCloudOptions {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        Ok(Self(ob))
+    }
+}
+
+impl OptRbCloudOptions {
+    pub fn extract_opt_cloud_options(
+        &self,
+        cloud_scheme: Option<CloudScheme>,
+        credential_provider: Option<Value>,
+    ) -> RbResult<Option<CloudOptions>> {
+        if self.0.is_nil() && credential_provider.is_none() {
+            return Ok(None);
+        }
+
+        let mut storage_options: Vec<(String, String)> = vec![];
+        let mut file_cache_ttl: u64 = 2;
+        let mut retry_config = CloudRetryConfig::default();
+
+        let storage_options_dict = Option::<RHash>::try_convert(self.0)?;
+
+        if let Some(storage_options_dict) = storage_options_dict {
+            storage_options.reserve(storage_options_dict.len());
+
+            storage_options_dict.foreach(|key: Value, value: Value| {
+                let key: String = key.funcall("to_s", ())?;
+
+                macro_rules! expected_type {
+                    ($key_name:expr, $type_name:expr) => {{
+                        |_| {
+                            let key_name = $key_name;
+                            let type_name = $type_name;
+                            RbValueError::new_err(format!(
+                                "invalid value for '{key_name}': '{value}' (expected {type_name})"
+                            ))
+                        }
+                    }};
+                }
+
+                match &*key {
+                    "file_cache_ttl" => {
+                        file_cache_ttl = u64::try_convert(value)
+                            .map_err(expected_type!("file_cache_ttl", "int"))?;
+                    }
+                    "max_retries" => {
+                        retry_config.max_retries = Option::<usize>::try_convert(value)
+                            .map_err(expected_type!("max_retries", "int"))?;
+                    }
+                    "retry_timeout_ms" => {
+                        retry_config.retry_timeout = Some(Duration::from_millis(
+                            u64::try_convert(value)
+                                .map_err(expected_type!("retry_timeout", "int"))?,
+                        ));
+                    }
+                    "retry_init_backoff_ms" => {
+                        retry_config.retry_init_backoff = Some(Duration::from_millis(
+                            u64::try_convert(value)
+                                .map_err(expected_type!("retry_init_backoff", "int"))?,
+                        ));
+                    }
+                    "retry_max_backoff_ms" => {
+                        retry_config.retry_max_backoff = Some(Duration::from_millis(
+                            u64::try_convert(value)
+                                .map_err(expected_type!("retry_max_backoff", "int"))?,
+                        ));
+                    }
+                    "retry_base_multiplier" => {
+                        retry_config.retry_base_multiplier = Some(TotalOrdWrap(
+                            f64::try_convert(value)
+                                .map_err(expected_type!("retry_base_multiplier", "float"))?,
+                        ));
+                    }
+                    _ => {
+                        let value =
+                            String::try_convert(value).map_err(expected_type!(&key, "str"))?;
+                        storage_options.push((key, value))
+                    }
+                }
+
+                Ok(ForEach::Continue)
+            })?;
+        }
+
+        let mut cloud_options = CloudOptions::from_untyped_config(cloud_scheme, storage_options)
+            .map_err(to_rb_err)?
+            .with_retry_config(retry_config);
+
+        if file_cache_ttl > 0 {
+            cloud_options.file_cache_ttl = file_cache_ttl;
+        }
+
+        Ok(Some(cloud_options))
+    }
+}

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

@@ -0,0 +1,4 @@
+pub mod cloud_options;
+pub mod scan_options;
+pub mod sink_options;
+pub mod sink_output;