Implement $write_csv() for DataFrame

NAMESPACE

@@ -151,6 +151,7 @@ S3method(tail,LazyFrame)
 S3method(unique,DataFrame)
 S3method(unique,LazyFrame)
 export(.pr)
+export(DataFrame_write_csv)
 export(as_polars_series)
 export(knit_print.DataFrame)
 export(pl)

R/dataframe__frame.R

@@ -1680,3 +1680,83 @@ 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.
+#' @export
+#' @rdname IO_write_csv
+#'
+#' @examples
+#' dat = pl$DataFrame(mtcars)
+#'
+
+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"
+) {
+
+  if (file_ext(path) != "csv") {
+    stop("Argument `path` must the path to a CSV file.")
+  }
+
+  if (length(quote_style) == 0 ||
+      !quote_style %in% c("always", "necessary", "non_numeric")) {
+    stop("Argument `quote_style` must be one of 'always', 'necessary', or 'non_numeric'.")
+  }
+
+  .pr$DataFrame$write_csv(
+    self,
+    path, has_header, utf8ToInt(separator), line_terminator, utf8ToInt(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,16 @@ 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_values, 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_values, quote_style)
+
 #' @export
 `$.DataFrame` <- function (self, name) { func <- DataFrame[[name]]; environment(func) <- environment(); func }
 

R/utils.R

@@ -639,3 +639,9 @@ is_bool = function(x) {
 dtypes_are_struct = function(dtypes) {
   sapply(dtypes, \(dt) pl$same_outer_dt(dt, pl$Struct()))
 }
+
+# from tools::file_ext()
+file_ext <- function(x) {
+  pos <- regexpr("\\.([[:alnum:]]+)$", x)
+  ifelse(pos > -1L, substring(x, pos + 1L), "")
+}

src/rust/src/rdataframe/mod.rs

@@ -11,6 +11,8 @@ use crate::rdatatype::RPolarsDataType;
 use crate::robj_to;
 use crate::rpolarserr::*;
 
+use polars::prelude::{CsvWriter, QuoteStyle, SerWriter};
+
 pub use lazy::dataframe::*;
 
 use crate::conversion_s_to_r::pl_series_to_list;
@@ -442,7 +444,57 @@ impl DataFrame {
             .map_err(polars_to_rpolars_err)
             .map(DataFrame)
     }
+
+    pub fn write_csv(
+        &mut 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 null = robj_to!(String, null_value).unwrap();
+        let path = robj_to!(str, path).unwrap();
+        let f = std::fs::File::create(path).unwrap();
+        let qs = parse_quote_style(quote_style);
+
+        CsvWriter::new(f)
+            .has_header(robj_to!(bool, has_header).unwrap())
+            .with_delimiter(robj_to!(u8, separator).unwrap())
+            .with_line_terminator(robj_to!(String, line_terminator).unwrap())
+            .with_quoting_char(robj_to!(u8, quote).unwrap())
+            .with_batch_size(robj_to!(usize, batch_size).unwrap())
+            .with_datetime_format(robj_to!(Option, String, datetime_format).unwrap())
+            .with_date_format(robj_to!(Option, String, date_format).unwrap())
+            .with_time_format(robj_to!(Option, String, time_format).unwrap())
+            .with_float_precision(robj_to!(Option, usize, float_precision).unwrap())
+            .with_null_value(null)
+            .with_quote_style(qs)
+            .finish(&mut self.0)
+            .map_err(polars_to_rpolars_err)
+    }
 }
+
+pub fn parse_quote_style(x: Robj) -> QuoteStyle {
+    match robj_to!(Option, String, x).unwrap_or_default().unwrap().as_str() {
+        "always" => QuoteStyle::Always,
+        "necessary" => QuoteStyle::Necessary,
+        "non_numeric" => QuoteStyle::NonNumeric,
+        // "never" is available in rust-polars devel only for now (will be added in 0.34)
+        // "never" => QuoteStyle::Never,
+        _ => panic!("polars internal error: `quote_style` must be 'always', 'necessary' or 'non_numeric'.")
+    }
+}
+
+
 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

tests/testthat/test-csv.R

@@ -22,3 +22,69 @@ test_that("csv read iris", {
     iris
   )
 })
+
+
+dat = mtcars
+dat[c(1, 3, 9, 12), c(3, 4, 5)] = NA
+dat_pl = pl$DataFrame(dat)
+temp_noext = tempfile()
+temp_out = tempfile(fileext = ".csv")
+
+test_that("write_csv: path works", {
+  expect_error(
+    dat_pl$write_csv(temp_noext),
+    "must the path to a CSV file"
+  )
+
+  dat_pl$write_csv(temp_out)
+  expect_identical(
+    pl$read_csv(temp_out)$to_data_frame(),
+    dat,
+    ignore_attr = TRUE # rownames are lost when writing / reading from CSV
+  )
+})
+
+test_that("write_csv: null_values works", {
+  expect_error(
+    dat_pl$write_csv(temp_out, null_values = NULL)
+  )
+  dat_pl$write_csv(temp_out, null_values = "hello")
+  tmp = pl$read_csv(temp_out)$to_data_frame()
+  expect_true(is.character(tmp$disp) && is.character(tmp$hp) && is.character(tmp$drat))
+  expect_equal(tmp[1:2, "disp"], c("hello", "160.0"))
+})
+
+
+test_that("write_csv: separator works", {
+  dat_pl$write_csv(temp_out, separator = "|")
+  expect_identical(
+    pl$read_csv(temp_out, sep = "|")$to_data_frame(),
+    dat,
+    ignore_attr = TRUE # rownames are lost when writing / reading from CSV
+  )
+})
+
+test_that("write_csv: quote_style and quote works", {
+  dat_pl2 = pl$DataFrame(iris)
+
+  expect_error(
+    dat_pl2$write_csv(temp_out, quote_style = "foo"),
+    "must be one of"
+  )
+
+  dat_pl2$write_csv(temp_out, quote_style = "always", quote = "+")
+  expect_identical(
+    head(pl$read_csv(temp_out)$to_data_frame()[["+Sepal.Length+"]], n = 2),
+    c("+5.1+", "+4.9+")
+  )
+
+  dat_pl2$write_csv(temp_out, quote_style = "non_numeric", quote = "+")
+  expect_identical(
+    head(pl$read_csv(temp_out)$to_data_frame()[["+Sepal.Length+"]], n = 2),
+    c(5.1, 4.9)
+  )
+  expect_identical(
+    head(pl$read_csv(temp_out)$to_data_frame()[["+Species+"]], n = 2),
+    c("+setosa+", "+setosa+")
+  )
+})