polars-df 0.3.1 → 0.5.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
@@ -115,7 +120,7 @@ impl RbDataFrame {
         let comment_char: Option<String> = arguments[17].try_convert()?;
         let quote_char: Option<String> = arguments[18].try_convert()?;
         let null_values: Option<Wrap<NullValues>> = arguments[19].try_convert()?;
-        let parse_dates: bool = arguments[20].try_convert()?;
+        let try_parse_dates: bool = arguments[20].try_convert()?;
         let skip_rows_after_header: usize = arguments[21].try_convert()?;
         let row_count: Option<(String, IdxSize)> = arguments[22].try_convert()?;
         let sample_size: usize = arguments[23].try_convert()?;
@@ -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| {
@@ -168,12 +175,12 @@ impl RbDataFrame {
             .with_columns(columns)
             .with_n_threads(n_threads)
             .with_path(path)
-            .with_dtypes(overwrite_dtype.as_ref())
+            .with_dtypes(overwrite_dtype.map(Arc::new))
             .with_dtypes_slice(overwrite_dtype_slice.as_deref())
             .low_memory(low_memory)
             .with_comment_char(comment_char)
             .with_null_values(null_values)
-            .with_parse_dates(parse_dates)
+            .with_try_parse_dates(try_parse_dates)
             .with_quote_char(quote_char)
             .with_end_of_line_char(eol_char)
             .with_skip_rows_after_header(skip_rows_after_header)
@@ -184,6 +191,7 @@ impl RbDataFrame {
         Ok(df.into())
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn read_parquet(
         rb_f: Value,
         columns: Option<Vec<String>>,
@@ -192,6 +200,8 @@ impl RbDataFrame {
         parallel: Wrap<ParallelStrategy>,
         row_count: 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 mmap_bytes_r = get_mmap_bytes_reader(rb_f)?;
@@ -202,6 +212,8 @@ impl RbDataFrame {
             .with_n_rows(n_rows)
             .with_row_count(row_count)
             .set_low_memory(low_memory)
+            .use_statistics(use_statistics)
+            .set_rechunk(rechunk)
             .finish()
             .map_err(RbPolarsErr::from)?;
         Ok(RbDataFrame::new(df))
@@ -254,7 +266,7 @@ impl RbDataFrame {
         use polars::io::avro::AvroWriter;
 
         if let Ok(s) = rb_f.try_convert::<String>() {
-            let f = std::fs::File::create(&s).unwrap();
+            let f = std::fs::File::create(s).unwrap();
             AvroWriter::new(f)
                 .with_compression(compression.0)
                 .finish(&mut self.df.borrow_mut())
@@ -339,7 +351,7 @@ impl RbDataFrame {
         // 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.clone();
+                *name = new_name.to_string();
             }
         }
         let rbdf = Self::finish_from_rows(
@@ -348,17 +360,19 @@ impl RbDataFrame {
             schema_overwrite.map(|wrap| wrap.0),
         )?;
 
-        rbdf.df
-            .borrow_mut()
-            .get_columns_mut()
-            .iter_mut()
-            .zip(&names)
-            .for_each(|(s, name)| {
-                s.rename(name);
-            });
+        unsafe {
+            rbdf.df
+                .borrow_mut()
+                .get_columns_mut()
+                .iter_mut()
+                .zip(&names)
+                .for_each(|(s, name)| {
+                    s.rename(name);
+                });
+        }
         let length = names.len();
         if names.into_iter().collect::<PlHashSet<_>>().len() != length {
-            let err = PolarsError::SchemaMisMatch("duplicate column names found".into());
+            let err = PolarsError::SchemaMismatch("duplicate column names found".into());
             Err(RbPolarsErr::from(err))?;
         }
 
@@ -394,7 +408,7 @@ impl RbDataFrame {
         let null = null_value.unwrap_or_default();
 
         if let Ok(s) = rb_f.try_convert::<String>() {
-            let f = std::fs::File::create(&s).unwrap();
+            let f = std::fs::File::create(s).unwrap();
             // no need for a buffered writer, because the csv writer does internal buffering
             CsvWriter::new(f)
                 .has_header(has_header)
@@ -436,7 +450,7 @@ impl RbDataFrame {
         compression: Wrap<Option<IpcCompression>>,
     ) -> RbResult<()> {
         if let Ok(s) = rb_f.try_convert::<String>() {
-            let f = std::fs::File::create(&s).unwrap();
+            let f = std::fs::File::create(s).unwrap();
             IpcWriter::new(f)
                 .with_compression(compression.0)
                 .finish(&mut self.df.borrow_mut())
@@ -505,12 +519,10 @@ impl RbDataFrame {
                 }
             }
         }
-        let st = st?;
+        let _st = st?;
 
-        match st {
-            // TODO
-            _ => None,
-        }
+        // TODO
+        None
     }
 
     pub fn write_parquet(
@@ -524,7 +536,7 @@ impl RbDataFrame {
         let compression = parse_parquet_compression(&compression, compression_level)?;
 
         if let Ok(s) = rb_f.try_convert::<String>() {
-            let f = std::fs::File::create(&s).unwrap();
+            let f = std::fs::File::create(s).unwrap();
             ParquetWriter::new(f)
                 .with_compression(compression)
                 .with_statistics(statistics)
@@ -627,7 +639,7 @@ impl RbDataFrame {
     }
 
     pub fn get_columns(&self) -> RArray {
-        let cols = self.df.borrow().get_columns().clone();
+        let cols = self.df.borrow().get_columns().to_vec();
         to_rbseries_collection(cols)
     }
 
@@ -881,10 +893,11 @@ impl RbDataFrame {
         variable_name: Option<String>,
     ) -> RbResult<Self> {
         let args = MeltArgs {
-            id_vars,
-            value_vars,
-            value_name,
-            variable_name,
+            id_vars: strings_to_smartstrings(id_vars),
+            value_vars: strings_to_smartstrings(value_vars),
+            value_name: value_name.map(|s| s.into()),
+            variable_name: variable_name.map(|s| s.into()),
+            streamable: false,
         };
 
         let df = self.df.borrow().melt2(args).map_err(RbPolarsErr::from)?;
@@ -897,22 +910,23 @@ impl RbDataFrame {
         values: Vec<String>,
         index: Vec<String>,
         columns: Vec<String>,
-        aggregate_expr: &RbExpr,
         maintain_order: bool,
         sort_columns: bool,
+        aggregate_expr: Option<&RbExpr>,
         separator: Option<String>,
     ) -> RbResult<Self> {
         let fun = match maintain_order {
             true => pivot_stable,
             false => pivot,
         };
+        let agg_expr = aggregate_expr.map(|aggregate_expr| aggregate_expr.inner.clone());
         let df = fun(
             &self.df.borrow(),
             values,
             index,
             columns,
-            aggregate_expr.inner.clone(),
             sort_columns,
+            agg_expr,
             separator.as_deref(),
         )
         .map_err(RbPolarsErr::from)?;
@@ -933,21 +947,6 @@ impl RbDataFrame {
         self.df.borrow().shift(periods).into()
     }
 
-    pub fn unique(
-        &self,
-        maintain_order: bool,
-        subset: Option<Vec<String>>,
-        keep: Wrap<UniqueKeepStrategy>,
-    ) -> RbResult<Self> {
-        let subset = subset.as_ref().map(|v| v.as_ref());
-        let df = match maintain_order {
-            true => self.df.borrow().unique_stable(subset, keep.0),
-            false => self.df.borrow().unique(subset, keep.0),
-        }
-        .map_err(RbPolarsErr::from)?;
-        Ok(df.into())
-    }
-
     pub fn lazy(&self) -> RbLazyFrame {
         self.df.borrow().clone().lazy().into()
     }

data/ext/polars/src/error.rs

@@ -22,10 +22,6 @@ impl RbPolarsErr {
     pub fn other(message: String) -> Error {
         Error::new(exception::runtime_error(), message)
     }
-
-    pub fn todo() -> Error {
-        Error::new(exception::runtime_error(), "not implemented yet")
-    }
 }
 
 pub struct RbValueError {}

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