polars-df 0.8.0 → 0.9.0

data/ext/polars/src/dataframe.rs

@@ -6,12 +6,13 @@ use polars::frame::row::{rows_to_schema_supertypes, Row};
 use polars::frame::NullStrategy;
 use polars::io::avro::AvroCompression;
 use polars::io::mmap::ReaderBytes;
-use polars::io::RowCount;
+use polars::io::RowIndex;
 use polars::prelude::pivot::{pivot, pivot_stable};
 use polars::prelude::*;
 use polars_core::utils::try_get_supertype;
 use std::cell::RefCell;
 use std::io::{BufWriter, Cursor};
+use std::num::NonZeroUsize;
 use std::ops::Deref;
 
 use crate::conversion::*;
@@ -45,44 +46,51 @@ impl RbDataFrame {
     fn finish_from_rows(
         rows: Vec<Row>,
         infer_schema_length: Option<usize>,
-        schema_overwrite: Option<Schema>,
+        schema: Option<Schema>,
+        schema_overrides_by_idx: Option<Vec<(usize, DataType)>>,
     ) -> RbResult<Self> {
-        // object builder must be registered.
+        // Object builder must be registered
         crate::on_startup::register_object_builder();
 
-        let schema =
+        let mut final_schema =
             rows_to_schema_supertypes(&rows, infer_schema_length.map(|n| std::cmp::max(1, n)))
                 .map_err(RbPolarsErr::from)?;
-        // replace inferred nulls with boolean
-        let fields = schema.iter_fields().map(|mut fld| match fld.data_type() {
-            DataType::Null => {
-                // fld.coerce(DataType::Boolean);
-                fld
-            }
-            DataType::Decimal(_, _) => {
-                fld.coerce(DataType::Decimal(None, None));
-                fld
+
+        // Erase scale from inferred decimals.
+        for dtype in final_schema.iter_dtypes_mut() {
+            if let DataType::Decimal(_, _) = dtype {
+                *dtype = DataType::Decimal(None, None)
             }
-            _ => fld,
-        });
-        let mut schema = Schema::from_iter(fields);
+        }
 
-        if let Some(schema_overwrite) = schema_overwrite {
-            for (i, (name, dtype)) in schema_overwrite.into_iter().enumerate() {
-                if let Some((name_, dtype_)) = schema.get_at_index_mut(i) {
+        // Integrate explicit/inferred schema.
+        if let Some(schema) = schema {
+            for (i, (name, dtype)) in schema.into_iter().enumerate() {
+                if let Some((name_, dtype_)) = final_schema.get_at_index_mut(i) {
                     *name_ = name;
 
-                    // if user sets dtype unknown, we use the inferred datatype
+                    // If schema dtype is Unknown, overwrite with inferred datatype.
                     if !matches!(dtype, DataType::Unknown) {
                         *dtype_ = dtype;
                     }
                 } else {
-                    schema.with_column(name, dtype);
+                    final_schema.with_column(name, dtype);
                 }
             }
         }
 
-        let df = DataFrame::from_rows_and_schema(&rows, &schema).map_err(RbPolarsErr::from)?;
+        // Optional per-field overrides; these supersede default/inferred dtypes.
+        if let Some(overrides) = schema_overrides_by_idx {
+            for (i, dtype) in overrides {
+                if let Some((_, dtype_)) = final_schema.get_at_index_mut(i) {
+                    if !matches!(dtype, DataType::Unknown) {
+                        *dtype_ = dtype;
+                    }
+                }
+            }
+        }
+        let df =
+            DataFrame::from_rows_and_schema(&rows, &final_schema).map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }
 
@@ -125,7 +133,7 @@ impl RbDataFrame {
         let null_values = Option::<Wrap<NullValues>>::try_convert(arguments[19])?;
         let try_parse_dates = bool::try_convert(arguments[20])?;
         let skip_rows_after_header = usize::try_convert(arguments[21])?;
-        let row_count = Option::<(String, IdxSize)>::try_convert(arguments[22])?;
+        let row_index = Option::<(String, IdxSize)>::try_convert(arguments[22])?;
         let sample_size = usize::try_convert(arguments[23])?;
         let eol_char = String::try_convert(arguments[24])?;
         // end arguments
@@ -133,7 +141,7 @@ impl RbDataFrame {
         let null_values = null_values.map(|w| w.0);
         let eol_char = eol_char.as_bytes()[0];
 
-        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
 
         let quote_char = if let Some(s) = quote_char {
             if s.is_empty() {
@@ -186,7 +194,7 @@ impl RbDataFrame {
             .with_quote_char(quote_char)
             .with_end_of_line_char(eol_char)
             .with_skip_rows_after_header(skip_rows_after_header)
-            .with_row_count(row_count)
+            .with_row_index(row_index)
             .sample_size(sample_size)
             .finish()
             .map_err(RbPolarsErr::from)?;
@@ -200,19 +208,19 @@ impl RbDataFrame {
         projection: Option<Vec<usize>>,
         n_rows: Option<usize>,
         parallel: Wrap<ParallelStrategy>,
-        row_count: Option<(String, IdxSize)>,
+        row_index: Option<(String, IdxSize)>,
         low_memory: bool,
         use_statistics: bool,
         rechunk: bool,
     ) -> RbResult<Self> {
-        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
         let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
         let df = ParquetReader::new(mmap_bytes_r)
             .with_projection(projection)
             .with_columns(columns)
             .read_parallel(parallel.0)
             .with_n_rows(n_rows)
-            .with_row_count(row_count)
+            .with_row_index(row_index)
             .set_low_memory(low_memory)
             .use_statistics(use_statistics)
             .set_rechunk(rechunk)
@@ -226,16 +234,16 @@ impl RbDataFrame {
         columns: Option<Vec<String>>,
         projection: Option<Vec<usize>>,
         n_rows: Option<usize>,
-        row_count: Option<(String, IdxSize)>,
+        row_index: Option<(String, IdxSize)>,
         memory_map: bool,
     ) -> RbResult<Self> {
-        let row_count = row_count.map(|(name, offset)| RowCount { name, offset });
+        let row_index = row_index.map(|(name, offset)| RowIndex { name, offset });
         let mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
         let df = IpcReader::new(mmap_bytes_r)
             .with_projection(projection)
             .with_columns(columns)
             .with_n_rows(n_rows)
-            .with_row_count(row_count)
+            .with_row_index(row_index)
             .memory_mapped(memory_map)
             .finish()
             .map_err(RbPolarsErr::from)?;
@@ -352,7 +360,7 @@ impl RbDataFrame {
     pub fn read_rows(
         rb_rows: RArray,
         infer_schema_length: Option<usize>,
-        schema_overwrite: Option<Wrap<Schema>>,
+        schema: Option<Wrap<Schema>>,
     ) -> RbResult<Self> {
         let mut rows = Vec::with_capacity(rb_rows.len());
         for v in rb_rows.each() {
@@ -363,30 +371,34 @@ impl RbDataFrame {
             }
             rows.push(Row(row));
         }
-        Self::finish_from_rows(
-            rows,
-            infer_schema_length,
-            schema_overwrite.map(|wrap| wrap.0),
-        )
+        Self::finish_from_rows(rows, infer_schema_length, schema.map(|wrap| wrap.0), None)
     }
 
     pub fn read_hashes(
         dicts: Value,
         infer_schema_length: Option<usize>,
-        schema_overwrite: Option<Wrap<Schema>>,
+        schema: Option<Wrap<Schema>>,
+        schema_overrides: Option<Wrap<Schema>>,
     ) -> RbResult<Self> {
-        let (rows, mut names) = dicts_to_rows(&dicts, infer_schema_length.unwrap_or(50))?;
+        let mut schema_columns = PlIndexSet::new();
+        if let Some(s) = &schema {
+            schema_columns.extend(s.0.iter_names().map(|n| n.to_string()))
+        }
+        let (rows, names) = dicts_to_rows(&dicts, infer_schema_length, schema_columns)?;
 
-        // ensure the new names are used
-        if let Some(schema) = &schema_overwrite {
-            for (new_name, name) in schema.0.iter_names().zip(names.iter_mut()) {
-                *name = new_name.to_string();
+        let mut schema_overrides_by_idx: Vec<(usize, DataType)> = Vec::new();
+        if let Some(overrides) = schema_overrides {
+            for (idx, name) in names.iter().enumerate() {
+                if let Some(dtype) = overrides.0.get(name) {
+                    schema_overrides_by_idx.push((idx, dtype.clone()));
+                }
             }
         }
         let rbdf = Self::finish_from_rows(
             rows,
             infer_schema_length,
-            schema_overwrite.map(|wrap| wrap.0),
+            schema.map(|wrap| wrap.0),
+            Some(schema_overrides_by_idx),
         )?;
 
         unsafe {
@@ -427,13 +439,14 @@ impl RbDataFrame {
         include_header: bool,
         separator: u8,
         quote_char: u8,
-        batch_size: usize,
+        batch_size: Wrap<NonZeroUsize>,
         datetime_format: Option<String>,
         date_format: Option<String>,
         time_format: Option<String>,
         float_precision: Option<usize>,
         null_value: Option<String>,
     ) -> RbResult<()> {
+        let batch_size = batch_size.0;
         let null = null_value.unwrap_or_default();
 
         if let Ok(s) = String::try_convert(rb_f) {
@@ -794,12 +807,11 @@ impl RbDataFrame {
         self.df.borrow().get_column_index(&name)
     }
 
-    // TODO remove clone
-    pub fn column(&self, name: String) -> RbResult<RbSeries> {
+    pub fn get_column(&self, name: String) -> RbResult<RbSeries> {
         self.df
             .borrow()
             .column(&name)
-            .map(|v| v.clone().into())
+            .map(|s| RbSeries::new(s.clone()))
             .map_err(RbPolarsErr::from)
     }
 
@@ -887,11 +899,11 @@ impl RbDataFrame {
         }
     }
 
-    pub fn with_row_count(&self, name: String, offset: Option<IdxSize>) -> RbResult<Self> {
+    pub fn with_row_index(&self, name: String, offset: Option<IdxSize>) -> RbResult<Self> {
         let df = self
             .df
             .borrow()
-            .with_row_count(&name, offset)
+            .with_row_index(&name, offset)
             .map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }
@@ -922,9 +934,9 @@ impl RbDataFrame {
     #[allow(clippy::too_many_arguments)]
     pub fn pivot_expr(
         &self,
-        values: Vec<String>,
         index: Vec<String>,
         columns: Vec<String>,
+        values: Option<Vec<String>>,
         maintain_order: bool,
         sort_columns: bool,
         aggregate_expr: Option<&RbExpr>,
@@ -937,9 +949,9 @@ impl RbDataFrame {
         let agg_expr = aggregate_expr.map(|aggregate_expr| aggregate_expr.inner.clone());
         let df = fun(
             &self.df.borrow(),
-            values,
             index,
             columns,
+            values,
             sort_columns,
             agg_expr,
             separator.as_deref(),
@@ -1121,7 +1133,7 @@ impl RbDataFrame {
         };
         Ok(self
             .df
-            .borrow()
+            .borrow_mut()
             .transpose(keep_names_as.as_deref(), new_col_names)
             .map_err(RbPolarsErr::from)?
             .into())
@@ -1163,4 +1175,8 @@ impl RbDataFrame {
         let df = self.df.borrow().unnest(names).map_err(RbPolarsErr::from)?;
         Ok(df.into())
     }
+
+    pub fn clear(&self) -> Self {
+        self.df.borrow().clear().into()
+    }
 }

data/ext/polars/src/expr/array.rs

@@ -1,3 +1,5 @@
+use polars::prelude::*;
+
 use crate::RbExpr;
 
 impl RbExpr {
@@ -12,4 +14,76 @@ impl RbExpr {
     pub fn array_sum(&self) -> Self {
         self.inner.clone().arr().sum().into()
     }
+
+    pub fn arr_unique(&self, maintain_order: bool) -> Self {
+        if maintain_order {
+            self.inner.clone().arr().unique_stable().into()
+        } else {
+            self.inner.clone().arr().unique().into()
+        }
+    }
+
+    pub fn arr_to_list(&self) -> Self {
+        self.inner.clone().arr().to_list().into()
+    }
+
+    pub fn arr_all(&self) -> Self {
+        self.inner.clone().arr().all().into()
+    }
+
+    pub fn arr_any(&self) -> Self {
+        self.inner.clone().arr().any().into()
+    }
+
+    pub fn arr_sort(&self, descending: bool, nulls_last: bool) -> Self {
+        self.inner
+            .clone()
+            .arr()
+            .sort(SortOptions {
+                descending,
+                nulls_last,
+                ..Default::default()
+            })
+            .into()
+    }
+
+    pub fn arr_reverse(&self) -> Self {
+        self.inner.clone().arr().reverse().into()
+    }
+
+    pub fn arr_arg_min(&self) -> Self {
+        self.inner.clone().arr().arg_min().into()
+    }
+
+    pub fn arr_arg_max(&self) -> Self {
+        self.inner.clone().arr().arg_max().into()
+    }
+
+    pub fn arr_get(&self, index: &RbExpr) -> Self {
+        self.inner.clone().arr().get(index.inner.clone()).into()
+    }
+
+    pub fn arr_join(&self, separator: &RbExpr, ignore_nulls: bool) -> Self {
+        self.inner
+            .clone()
+            .arr()
+            .join(separator.inner.clone(), ignore_nulls)
+            .into()
+    }
+
+    pub fn arr_contains(&self, other: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .arr()
+            .contains(other.inner.clone())
+            .into()
+    }
+
+    pub fn arr_count_matches(&self, expr: &RbExpr) -> Self {
+        self.inner
+            .clone()
+            .arr()
+            .count_matches(expr.inner.clone())
+            .into()
+    }
 }

data/ext/polars/src/expr/datetime.rs

@@ -61,6 +61,14 @@ impl RbExpr {
         self.inner.clone().dt().month_end().into()
     }
 
+    pub fn dt_base_utc_offset(&self) -> Self {
+        self.inner.clone().dt().base_utc_offset().into()
+    }
+
+    pub fn dt_dst_offset(&self) -> Self {
+        self.inner.clone().dt().dst_offset().into()
+    }
+
     pub fn dt_round(&self, every: String, offset: String) -> Self {
         self.inner.clone().dt().round(&every, &offset).into()
     }
@@ -149,73 +157,31 @@ impl RbExpr {
         self.inner.clone().dt().timestamp(tu.0).into()
     }
 
-    pub fn duration_days(&self) -> Self {
-        self.inner
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.days().into_series())),
-                GetOutput::from_type(DataType::Int64),
-            )
-            .into()
+    pub fn dt_total_days(&self) -> Self {
+        self.inner.clone().dt().total_days().into()
     }
 
-    pub fn duration_hours(&self) -> Self {
-        self.inner
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.hours().into_series())),
-                GetOutput::from_type(DataType::Int64),
-            )
-            .into()
+    pub fn dt_total_hours(&self) -> Self {
+        self.inner.clone().dt().total_hours().into()
     }
 
-    pub fn duration_minutes(&self) -> Self {
-        self.inner
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.minutes().into_series())),
-                GetOutput::from_type(DataType::Int64),
-            )
-            .into()
+    pub fn dt_total_minutes(&self) -> Self {
+        self.inner.clone().dt().total_minutes().into()
     }
 
-    pub fn duration_seconds(&self) -> Self {
-        self.inner
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.seconds().into_series())),
-                GetOutput::from_type(DataType::Int64),
-            )
-            .into()
+    pub fn dt_total_seconds(&self) -> Self {
+        self.inner.clone().dt().total_seconds().into()
     }
 
-    pub fn duration_milliseconds(&self) -> Self {
-        self.inner
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.milliseconds().into_series())),
-                GetOutput::from_type(DataType::Int64),
-            )
-            .into()
+    pub fn dt_total_milliseconds(&self) -> Self {
+        self.inner.clone().dt().total_milliseconds().into()
     }
 
-    pub fn duration_microseconds(&self) -> Self {
-        self.inner
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.microseconds().into_series())),
-                GetOutput::from_type(DataType::Int64),
-            )
-            .into()
+    pub fn dt_total_microseconds(&self) -> Self {
+        self.inner.clone().dt().total_microseconds().into()
     }
 
-    pub fn duration_nanoseconds(&self) -> Self {
-        self.inner
-            .clone()
-            .map(
-                |s| Ok(Some(s.duration()?.nanoseconds().into_series())),
-                GetOutput::from_type(DataType::Int64),
-            )
-            .into()
+    pub fn dt_total_nanoseconds(&self) -> Self {
+        self.inner.clone().dt().total_nanoseconds().into()
     }
 }