polars-df 0.4.0 → 0.6.0

data/ext/polars/src/dataframe.rs

@@ -48,20 +48,25 @@ impl RbDataFrame {
         crate::object::register_object_builder();
 
         let schema =
-            rows_to_schema_supertypes(&rows, infer_schema_length).map_err(RbPolarsErr::from)?;
+            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
+            }
             _ => fld,
         });
-        let mut schema = Schema::from(fields);
+        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_index_mut(i) {
+                if let Some((name_, dtype_)) = schema.get_at_index_mut(i) {
                     *name_ = name;
 
                     // if user sets dtype unknown, we use the inferred datatype
@@ -139,11 +144,13 @@ impl RbDataFrame {
         };
 
         let overwrite_dtype = overwrite_dtype.map(|overwrite_dtype| {
-            let fields = overwrite_dtype.iter().map(|(name, dtype)| {
-                let dtype = dtype.0.clone();
-                Field::new(name, dtype)
-            });
-            Schema::from(fields)
+            overwrite_dtype
+                .iter()
+                .map(|(name, dtype)| {
+                    let dtype = dtype.0.clone();
+                    Field::new(name, dtype)
+                })
+                .collect::<Schema>()
         });
 
         let overwrite_dtype_slice = overwrite_dtype_slice.map(|overwrite_dtype| {
@@ -449,12 +456,14 @@ impl RbDataFrame {
                 .finish(&mut self.df.borrow_mut())
                 .map_err(RbPolarsErr::from)?;
         } else {
-            let mut buf = get_file_like(rb_f, true)?;
-
+            let mut buf = Cursor::new(Vec::new());
             IpcWriter::new(&mut buf)
                 .with_compression(compression.0)
                 .finish(&mut self.df.borrow_mut())
                 .map_err(RbPolarsErr::from)?;
+            // TODO less copying
+            let rb_str = RString::from_slice(&buf.into_inner());
+            rb_f.funcall::<_, _, Value>("write", (rb_str,))?;
         }
         Ok(())
     }
@@ -512,12 +521,10 @@ impl RbDataFrame {
                 }
             }
         }
-        let st = st?;
+        let _st = st?;
 
-        match st {
-            // TODO
-            _ => None,
-        }
+        // TODO
+        None
     }
 
     pub fn write_parquet(
@@ -789,22 +796,6 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
-    pub fn sort(&self, by_column: String, reverse: bool, nulls_last: bool) -> RbResult<Self> {
-        let df = self
-            .df
-            .borrow()
-            .sort_with_options(
-                &by_column,
-                SortOptions {
-                    descending: reverse,
-                    nulls_last,
-                    multithreaded: true,
-                },
-            )
-            .map_err(RbPolarsErr::from)?;
-        Ok(RbDataFrame::new(df))
-    }
-
     pub fn replace(&self, column: String, new_col: &RbSeries) -> RbResult<()> {
         self.df
             .borrow_mut()
@@ -914,10 +905,7 @@ impl RbDataFrame {
             true => pivot_stable,
             false => pivot,
         };
-        let agg_expr = match aggregate_expr {
-            Some(aggregate_expr) => Some(aggregate_expr.inner.clone()),
-            None => None,
-        };
+        let agg_expr = aggregate_expr.map(|aggregate_expr| aggregate_expr.inner.clone());
         let df = fun(
             &self.df.borrow(),
             values,
@@ -931,11 +919,16 @@ impl RbDataFrame {
         Ok(RbDataFrame::new(df))
     }
 
-    pub fn partition_by(&self, groups: Vec<String>, stable: bool) -> RbResult<RArray> {
-        let out = if stable {
-            self.df.borrow().partition_by_stable(groups)
+    pub fn partition_by(
+        &self,
+        by: Vec<String>,
+        maintain_order: bool,
+        include_key: bool,
+    ) -> RbResult<RArray> {
+        let out = if maintain_order {
+            self.df.borrow().partition_by_stable(by, include_key)
         } else {
-            self.df.borrow().partition_by(groups)
+            self.df.borrow().partition_by(by, include_key)
         }
         .map_err(RbPolarsErr::from)?;
         Ok(RArray::from_iter(out.into_iter().map(RbDataFrame::new)))
@@ -1022,13 +1015,18 @@ impl RbDataFrame {
         &self,
         columns: Option<Vec<String>>,
         separator: Option<String>,
+        drop_first: bool,
     ) -> RbResult<Self> {
         let df = match columns {
             Some(cols) => self.df.borrow().columns_to_dummies(
                 cols.iter().map(|x| x as &str).collect(),
                 separator.as_deref(),
+                drop_first,
             ),
-            None => self.df.borrow().to_dummies(separator.as_deref()),
+            None => self
+                .df
+                .borrow()
+                .to_dummies(separator.as_deref(), drop_first),
         }
         .map_err(RbPolarsErr::from)?;
         Ok(df.into())

data/ext/polars/src/error.rs

@@ -22,9 +22,13 @@ impl RbPolarsErr {
     pub fn other(message: String) -> Error {
         Error::new(exception::runtime_error(), message)
     }
+}
+
+pub struct RbTypeError {}
 
-    pub fn todo() -> Error {
-        Error::new(exception::runtime_error(), "not implemented yet")
+impl RbTypeError {
+    pub fn new_err(message: String) -> Error {
+        Error::new(exception::type_error(), message)
     }
 }
 

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

@@ -0,0 +1,15 @@
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn array_max(&self) -> Self {
+        self.inner.clone().arr().max().into()
+    }
+
+    pub fn array_min(&self) -> Self {
+        self.inner.clone().arr().min().into()
+    }
+
+    pub fn array_sum(&self) -> Self {
+        self.inner.clone().arr().sum().into()
+    }
+}

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

@@ -0,0 +1,69 @@
+use polars::prelude::*;
+
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn bin_contains(&self, lit: Vec<u8>) -> Self {
+        self.inner.clone().binary().contains_literal(lit).into()
+    }
+
+    pub fn bin_ends_with(&self, sub: Vec<u8>) -> Self {
+        self.inner.clone().binary().ends_with(sub).into()
+    }
+
+    pub fn bin_starts_with(&self, sub: Vec<u8>) -> Self {
+        self.inner.clone().binary().starts_with(sub).into()
+    }
+
+    pub fn bin_hex_decode(&self, strict: bool) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| {
+                    s.binary()?
+                        .hex_decode(strict)
+                        .map(|s| Some(s.into_series()))
+                },
+                GetOutput::same_type(),
+            )
+            .with_fmt("bin.hex_decode")
+            .into()
+    }
+
+    pub fn bin_base64_decode(&self, strict: bool) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| {
+                    s.binary()?
+                        .base64_decode(strict)
+                        .map(|s| Some(s.into_series()))
+                },
+                GetOutput::same_type(),
+            )
+            .with_fmt("bin.base64_decode")
+            .into()
+    }
+
+    pub fn bin_hex_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.binary().map(|s| Some(s.hex_encode().into_series())),
+                GetOutput::same_type(),
+            )
+            .with_fmt("bin.hex_encode")
+            .into()
+    }
+
+    pub fn bin_base64_encode(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                move |s| s.binary().map(|s| Some(s.base64_encode().into_series())),
+                GetOutput::same_type(),
+            )
+            .with_fmt("bin.base64_encode")
+            .into()
+    }
+}

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

@@ -0,0 +1,10 @@
+use polars::prelude::*;
+
+use crate::conversion::Wrap;
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn cat_set_ordering(&self, ordering: Wrap<CategoricalOrdering>) -> Self {
+        self.inner.clone().cat().set_ordering(ordering.0).into()
+    }
+}

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

@@ -0,0 +1,223 @@
+use polars::prelude::*;
+
+use crate::conversion::Wrap;
+use crate::RbExpr;
+
+impl RbExpr {
+    pub fn dt_to_string(&self, format: String) -> Self {
+        self.inner.clone().dt().to_string(&format).into()
+    }
+
+    pub fn dt_offset_by(&self, by: String) -> Self {
+        let by = Duration::parse(&by);
+        self.inner.clone().dt().offset_by(by).into()
+    }
+
+    pub fn dt_epoch_seconds(&self) -> Self {
+        self.clone()
+            .inner
+            .map(
+                |s| {
+                    s.timestamp(TimeUnit::Milliseconds)
+                        .map(|ca| Some((ca / 1000).into_series()))
+                },
+                GetOutput::from_type(DataType::Int64),
+            )
+            .into()
+    }
+
+    pub fn dt_with_time_unit(&self, tu: Wrap<TimeUnit>) -> Self {
+        self.inner.clone().dt().with_time_unit(tu.0).into()
+    }
+
+    pub fn dt_convert_time_zone(&self, tz: TimeZone) -> Self {
+        self.inner.clone().dt().convert_time_zone(tz).into()
+    }
+
+    pub fn dt_cast_time_unit(&self, tu: Wrap<TimeUnit>) -> Self {
+        self.inner.clone().dt().cast_time_unit(tu.0).into()
+    }
+
+    pub fn dt_replace_time_zone(&self, tz: Option<String>, use_earliest: Option<bool>) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .replace_time_zone(tz, use_earliest)
+            .into()
+    }
+
+    #[allow(deprecated)]
+    pub fn dt_tz_localize(&self, tz: String) -> Self {
+        self.inner.clone().dt().tz_localize(tz).into()
+    }
+
+    pub fn dt_truncate(&self, every: String, offset: String) -> Self {
+        self.inner.clone().dt().truncate(&every, &offset).into()
+    }
+
+    pub fn dt_month_start(&self) -> Self {
+        self.inner.clone().dt().month_start().into()
+    }
+
+    pub fn dt_month_end(&self) -> Self {
+        self.inner.clone().dt().month_end().into()
+    }
+
+    pub fn dt_round(&self, every: String, offset: String) -> Self {
+        self.inner.clone().dt().round(&every, &offset).into()
+    }
+
+    pub fn dt_combine(&self, time: &Self, time_unit: Wrap<TimeUnit>) -> Self {
+        self.inner
+            .clone()
+            .dt()
+            .combine(time.inner.clone(), time_unit.0)
+            .into()
+    }
+
+    pub fn dt_year(&self) -> Self {
+        self.clone().inner.dt().year().into()
+    }
+
+    pub fn dt_is_leap_year(&self) -> Self {
+        self.clone().inner.dt().is_leap_year().into()
+    }
+
+    pub fn dt_iso_year(&self) -> Self {
+        self.clone().inner.dt().iso_year().into()
+    }
+
+    pub fn dt_quarter(&self) -> Self {
+        self.clone().inner.dt().quarter().into()
+    }
+
+    pub fn dt_month(&self) -> Self {
+        self.clone().inner.dt().month().into()
+    }
+
+    pub fn dt_week(&self) -> Self {
+        self.clone().inner.dt().week().into()
+    }
+
+    pub fn dt_weekday(&self) -> Self {
+        self.clone().inner.dt().weekday().into()
+    }
+
+    pub fn dt_day(&self) -> Self {
+        self.clone().inner.dt().day().into()
+    }
+
+    pub fn dt_ordinal_day(&self) -> Self {
+        self.clone().inner.dt().ordinal_day().into()
+    }
+
+    pub fn dt_time(&self) -> Self {
+        self.clone().inner.dt().time().into()
+    }
+
+    pub fn dt_date(&self) -> Self {
+        self.clone().inner.dt().date().into()
+    }
+
+    pub fn dt_datetime(&self) -> Self {
+        self.clone().inner.dt().datetime().into()
+    }
+
+    pub fn dt_hour(&self) -> Self {
+        self.clone().inner.dt().hour().into()
+    }
+
+    pub fn dt_minute(&self) -> Self {
+        self.clone().inner.dt().minute().into()
+    }
+
+    pub fn dt_second(&self) -> Self {
+        self.clone().inner.dt().second().into()
+    }
+
+    pub fn dt_millisecond(&self) -> Self {
+        self.clone().inner.dt().millisecond().into()
+    }
+
+    pub fn dt_microsecond(&self) -> Self {
+        self.clone().inner.dt().microsecond().into()
+    }
+
+    pub fn dt_nanosecond(&self) -> Self {
+        self.clone().inner.dt().nanosecond().into()
+    }
+
+    pub fn dt_timestamp(&self, tu: Wrap<TimeUnit>) -> Self {
+        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 duration_hours(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.hours().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .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 duration_seconds(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.seconds().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .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 duration_microseconds(&self) -> Self {
+        self.inner
+            .clone()
+            .map(
+                |s| Ok(Some(s.duration()?.microseconds().into_series())),
+                GetOutput::from_type(DataType::Int64),
+            )
+            .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()
+    }
+}