Implement $write_csv() for DataFrame

NEWS.md

@@ -44,6 +44,7 @@
 - New function `pl$raw_list` and class `rpolars_raw_list` a list of R Raw's, where missing is
   encoded as `NULL` to aid conversion to polars binary Series. Support back and forth conversion
   from polars binary literal and Series to R raw (#417).
+- New method `$write_csv()` for `DataFrame` (#414).
 - New method `$dt$time()` to extract the time from a `datetime` variable (#428).
 
 # polars 0.8.1

R/dataframe__frame.R

@@ -1680,3 +1680,75 @@ DataFrame_sample = function(
   ) |>
     unwrap("in $sample():")
 }
+
+
+
+#' Write to comma-separated values (CSV) file
+#'
+#' @param path File path to which the result should be written.
+#' @param has_header Whether to include header in the CSV output.
+#' @param separator Separate CSV fields with this symbol.
+#' @param line_terminator String used to end each row.
+#' @param quote Byte to use as quoting character.
+#' @param batch_size Number of rows that will be processed per thread.
+#' @param datetime_format A format string, with the specifiers defined by the
+#' chrono Rust crate. If no format specified, the default fractional-second
+#' precision is inferred from the maximum timeunit found in the frame’s Datetime
+#'  cols (if any).
+#' @param date_format A format string, with the specifiers defined by the chrono
+#' Rust crate.
+#' @param time_format A format string, with the specifiers defined by the chrono
+#' Rust crate.
+#' @param float_precision Number of decimal places to write, applied to both
+#' Float32 and Float64 datatypes.
+#' @param null_values A string representing null values (defaulting to the empty
+#' string).
+#' @param quote_style Determines the quoting strategy used.
+#' * `"necessary"` (default): This puts quotes around fields only when necessary.
+#'   They are necessary when fields contain a quote, delimiter or record
+#'   terminator. Quotes are also necessary when writing an empty record (which
+#'   is indistinguishable from a record with one empty field). This is the
+#'   default.
+#' * `"always"`: This puts quotes around every field.
+#' * `"non_numeric"`: This puts quotes around all fields that are non-numeric.
+#'   Namely, when writing a field that does not parse as a valid float or integer,
+#'   then quotes will be used even if they aren`t strictly necessary.
+
+# TODO: include "never" when bumping rust-polars to 0.34
+# * `"never"`: This never puts quotes around fields, even if that results in
+#   invalid CSV data (e.g.: by not quoting strings containing the separator).
+
+#' @return
+#' This doesn't return anything but creates a CSV file.
+#'
+#' @rdname IO_write_csv
+#'
+#' @examples
+#' dat = pl$DataFrame(mtcars)
+#'
+#' destination = tempfile(fileext = ".csv")
+#' dat$select(pl$col("drat", "mpg"))$write_csv(destination)
+#'
+#' pl$read_csv(destination)
+DataFrame_write_csv = function(
+    path,
+    has_header = TRUE,
+    separator = ",",
+    line_terminator = "\n",
+    quote = '"',
+    batch_size = 1024,
+    datetime_format = NULL,
+    date_format = NULL,
+    time_format = NULL,
+    float_precision = NULL,
+    null_values = "",
+    quote_style = "necessary") {
+  .pr$DataFrame$write_csv(
+    self,
+    path, has_header, separator, line_terminator, quote, batch_size,
+    datetime_format, date_format, time_format, float_precision,
+    null_values, quote_style
+  ) |>
+    unwrap("in $write_csv():") |>
+    invisible()
+}

R/extendr-wrappers.R

@@ -189,6 +189,8 @@ DataFrame$sample_n <- function(n, with_replacement, shuffle, seed) .Call(wrap__D
 
 DataFrame$sample_frac <- function(frac, with_replacement, shuffle, seed) .Call(wrap__DataFrame__sample_frac, self, frac, with_replacement, shuffle, seed)
 
+DataFrame$write_csv <- function(path, has_header, separator, line_terminator, quote, batch_size, datetime_format, date_format, time_format, float_precision, null_value, quote_style) .Call(wrap__DataFrame__write_csv, self, path, has_header, separator, line_terminator, quote, batch_size, datetime_format, date_format, time_format, float_precision, null_value, quote_style)
+
 #' @export
 `$.DataFrame` <- function (self, name) { func <- DataFrame[[name]]; environment(func) <- environment(); func }
 

src/rust/src/lazy/dsl.rs

@@ -1314,12 +1314,7 @@ impl Expr {
     }
 
     pub fn dt_time(&self) -> RResult<Self> {
-        Ok(self
-            .0
-            .clone()
-            .dt()
-            .time()
-            .into())
+        Ok(self.0.clone().dt().time().into())
     }
 
     pub fn dt_combine(&self, time: Robj, tu: Robj) -> RResult<Expr> {

src/rust/src/rdataframe/mod.rs

@@ -1,5 +1,5 @@
 use extendr_api::{extendr, prelude::*, rprintln, Rinternals};
-use polars::prelude::{self as pl, IntoLazy};
+use polars::prelude::{self as pl, IntoLazy, SerWriter};
 use std::result::Result;
 pub mod read_csv;
 pub mod read_ipc;
@@ -442,7 +442,41 @@ impl DataFrame {
             .map_err(polars_to_rpolars_err)
             .map(DataFrame)
     }
+
+    pub fn write_csv(
+        &self,
+        path: Robj,
+        has_header: Robj,
+        separator: Robj,
+        line_terminator: Robj,
+        quote: Robj,
+        batch_size: Robj,
+        datetime_format: Robj,
+        date_format: Robj,
+        time_format: Robj,
+        float_precision: Robj,
+        null_value: Robj,
+        quote_style: Robj,
+    ) -> RResult<()> {
+        let path = robj_to!(str, path)?;
+        let f = std::fs::File::create(path)?;
+        pl::CsvWriter::new(f)
+            .has_header(robj_to!(bool, has_header)?)
+            .with_delimiter(robj_to!(Utf8Byte, separator)?)
+            .with_line_terminator(robj_to!(String, line_terminator)?)
+            .with_quoting_char(robj_to!(Utf8Byte, quote)?)
+            .with_batch_size(robj_to!(usize, batch_size)?)
+            .with_datetime_format(robj_to!(Option, String, datetime_format)?)
+            .with_date_format(robj_to!(Option, String, date_format)?)
+            .with_time_format(robj_to!(Option, String, time_format)?)
+            .with_float_precision(robj_to!(Option, usize, float_precision)?)
+            .with_null_value(robj_to!(String, null_value)?)
+            .with_quote_style(robj_to!(QuoteStyle, quote_style)?)
+            .finish(&mut self.0.clone())
+            .map_err(polars_to_rpolars_err)
+    }
 }
+
 impl DataFrame {
     pub fn to_list_result(&self) -> Result<Robj, pl::PolarsError> {
         //convert DataFrame to Result of to R vectors, error if DataType is not supported

src/rust/src/utils/mod.rs

@@ -535,6 +535,29 @@ pub fn robj_to_usize(robj: extendr_api::Robj) -> RResult<usize> {
     robj_to_u64(robj).and_then(try_u64_into_usize)
 }
 
+pub fn robj_to_utf8_byte(robj: extendr_api::Robj) -> RResult<u8> {
+    let mut utf8_byte_iter = robj_to_str(robj)?.as_bytes().iter();
+    match (utf8_byte_iter.next(), utf8_byte_iter.next()) {
+        (Some(s), None) => Ok(*s),
+        (None, None) => rerr().plain("cannot extract single byte from empty string"),
+        (Some(_), Some(_)) => rerr().plain("multi byte-string not allowed"),
+        (None, Some(_)) => unreachable!("the iter() cannot yield Some after None(depleted)"),
+    }
+}
+
+pub fn robj_to_quote_style(robj: Robj) -> RResult<pl::QuoteStyle> {
+    match robj_to_str(robj.clone())? {
+        "always" => Ok(pl::QuoteStyle::Always),
+        "necessary" => Ok(pl::QuoteStyle::Necessary),
+        "non_numeric" => Ok(pl::QuoteStyle::NonNumeric),
+        // "never" is available in rust-polars devel only for now (will be added in 0.34)
+        // "never" => Ok(QuoteStyle::Never),
+        _ => rerr()
+            .plain("a `quote_style` must be 'always', 'necessary' or 'non_numeric'.")
+            .bad_robj(&robj),
+    }
+}
+
 fn err_no_nan<T>() -> RResult<T> {
     rerr().plain("any NA value is not allowed here".to_string())
 }
@@ -885,6 +908,10 @@ macro_rules! robj_to_inner {
         $crate::utils::robj_to_u8($a)
     };
 
+    (Utf8Byte, $a:ident) => {
+        $crate::utils::robj_to_utf8_byte($a)
+    };
+
     (char, $a:ident) => {
         $crate::utils::robj_to_char($a)
     };
@@ -985,6 +1012,10 @@ macro_rules! robj_to_inner {
         $crate::utils::robj_to_dataframe($a).map(|lf| lf.0)
     };
 
+    (QuoteStyle, $a:ident) => {
+        $crate::utils::robj_to_quote_style($a)
+    };
+
     (RArrow_schema, $a:ident) => {
         $crate::utils::robj_to_rarrow_schema($a)
     };