feat: Add partition_by(), deprecate partition_by_key() and partition_by_max_size() (#299)

Author: etiennebacher

.github/workflows/R-CMD-check.yml

@@ -55,7 +55,8 @@ jobs:
       - uses: r-lib/actions/setup-r@v2
         with:
           use-public-rspm: true
-          extra-repositories: 'https://community.r-multiverse.org'
+          extra-repositories: 'https://rpolars.r-universe.org'
+          # extra-repositories: 'https://community.r-multiverse.org'
 
       - name: Install R dependencies
         uses: r-lib/actions/setup-r-dependencies@v2

.github/workflows/covr.yml

@@ -42,7 +42,8 @@ jobs:
       - uses: r-lib/actions/setup-r@v2
         with:
           use-public-rspm: true
-          extra-repositories: 'https://community.r-multiverse.org'
+          extra-repositories: 'https://rpolars.r-universe.org'
+          # extra-repositories: 'https://community.r-multiverse.org'
 
       - uses: r-lib/actions/setup-r-dependencies@v2
         with:

NAMESPACE

@@ -94,6 +94,7 @@ export(is_polars_df)
 export(is_polars_expr)
 export(is_polars_lf)
 export(make_unique_id)
+export(partition_by)
 export(partition_by_key)
 export(partition_by_max_size)
 export(read_csv_polars)

NEWS.md

@@ -37,6 +37,13 @@
 * New argument `mkdir` in `write_parquet_polars()` (this already existed in
   `sink_parquet()`). (#298)
 
+* New (experimental) function `partition_by()` to write partitioned output in
+  `sink_*()` and `write_*_polars()`. The following functions are deprecated and
+  will be removed in a future release (#299):
+
+  - `partition_by_key()` can be replaced with `partition_by(key =)`
+  - `partition_by_max_size()` can be replaced with `partition_by(max_rows_per_file =)`
+
 ## Changes
 
 * `collect()` now returns a `tibble` instead of a `data.frame`, for consistency

R/partitioned-output.R

@@ -0,0 +1,116 @@
+#' Helper functions to export data as a partitioned output
+#'
+#' @description
+#' `r lifecycle::badge("experimental")`
+#'
+#' Partitioning schemes are used to write multiple files with `sink_*()` and
+#' `write_*_polars()` functions.
+#'
+#' - `partition_by()`: Configuration for writing to multiple output files.
+#'   Supports partitioning by key, file size limits, or both.
+#'
+#' The following functions are deprecated and will be removed in a future release:
+#'
+#' - `r lifecycle::badge("deprecated")` `partition_by_key()`: use
+#'   `partition_by(key = ...)` instead.
+#' - `r lifecycle::badge("deprecated")` `partition_by_max_size()`: use
+#'   `partition_by(max_rows_per_file = ...)` instead.
+#'
+#' @inheritParams rlang::args_dots_empty
+#' @param base_path The base path for the output files. Use the `mkdir` option
+#' of the `sink_*` or `write_*_polars()` functions to ensure directories in the
+#' path are created.
+#' @param key Something can be coerced to a list of Polars expressions. Used to
+#' partition by.
+#' @param include_key A bool indicating whether to include the key columns in
+#' the output files. Can only be used if `key` is specified, otherwise should
+#' be `NULL`.
+#' @param max_rows_per_file An integer-ish value indicating the maximum size in
+#' rows of each of the generated files.
+#' @param approximate_bytes_per_file An integer-ish value indicating approximate
+#' number of bytes to write to each file, or `NULL`. This is measured as the
+#' estimated size of the DataFrame in memory. Defaults to approximately 4GB when
+#' `key` is specified without `max_rows_per_file`, otherwise unlimited.
+#' @param by `r lifecycle::badge("deprecated")` Something can be coerced to a
+#' list of Polars expressions. Used to partition by. Use the `key` property of
+#' `partition_by()` instead.
+#' @param per_partition_sort_by `r lifecycle::badge("deprecated")` Something
+#' that can be coerced to a list of Polars expressions, or `NULL` (default).
+#' Used to sort over within each partition. Use the `per_partition_sort_by`
+#' property of `partition_by()` instead.
+#' @param max_size `r lifecycle::badge("deprecated")` An integer-ish value
+#' indicating the maximum size in rows of each of the generated files. Use the
+#' `max_rows_per_file` property of `partition_by()` instead.
+#'
+#' @name partitioned_output
+#' @export
+partition_by <- function(
+  base_path,
+  ...,
+  key = NULL,
+  include_key = NULL,
+  max_rows_per_file = NULL,
+  approximate_bytes_per_file = NULL
+) {
+  check_dots_empty()
+  pl$PartitionBy(
+    base_path = base_path,
+    key = key,
+    include_key = include_key,
+    max_rows_per_file = max_rows_per_file,
+    approximate_bytes_per_file = approximate_bytes_per_file
+  )
+}
+
+#' @name partitioned_output
+#' @export
+partition_by_key <- function(
+  base_path,
+  ...,
+  by,
+  include_key = TRUE,
+  per_partition_sort_by = NULL
+) {
+  check_dots_empty()
+
+  lifecycle::deprecate_warn(
+    when = "0.16.0",
+    what = "partition_by_key()",
+    details = "Please use `partition_by(key = )` instead.",
+    always = TRUE,
+  )
+
+  suppressWarnings(
+    pl$PartitionByKey(
+      base_path = base_path,
+      by = by,
+      include_key = include_key,
+      per_partition_sort_by = per_partition_sort_by
+    )
+  )
+}
+
+#' @rdname partitioned_output
+#' @export
+partition_by_max_size <- function(
+  base_path,
+  ...,
+  max_size,
+  per_partition_sort_by = NULL
+) {
+  check_dots_empty()
+
+  lifecycle::deprecate_warn(
+    when = "0.16.0",
+    what = "partition_by_max_size()",
+    details = "Please use `partition_by(max_rows_per_file = )` instead.",
+    always = TRUE,
+  )
+  suppressWarnings(
+    pl$PartitionMaxSize(
+      base_path = base_path,
+      max_size = max_size,
+      per_partition_sort_by = per_partition_sort_by
+    )
+  )
+}

R/sink.R

@@ -57,33 +57,9 @@
 #'
 #' ## Partitioned output
 #'
-#' It is possible to export a LazyFrame to multiple files, also called
-#' *partitioned output*. A partition can be determined in several ways:
-#'
-#' - by key(s): split by the values of keys. The amount of files that can be
-#'   written is not limited. However, when writing beyond a certain amount of
-#'   files, the data for the remaining partitions is buffered before writing to
-#'   the file.
-#' - by maximum number of rows: if the number of rows in a file reaches the
-#'   maximum number of rows, the file is closed and a new file is opened.
-
-# TODO: add this back when https://github.com/pola-rs/r-polars/issues/1522 is
-# solved
-# - by "sorted partition": this is a specialized version of partitioning by
-#   key. Whereas partitioning by key accepts data in any order, this scheme
-#   expects the input data to be pre-grouped or pre-sorted. This scheme suffers
-#   a lot less overhead, but may not be always applicable. Each new value of
-#   the key expressions starts a new partition, therefore repeating the same
-#   value multiple times may overwrite previous partitions.
-# These partitioning schemes can be used with the functions `partition_by_key()`,
-# `partition_by_max_size()`, and `partition_parted()`. See Examples below.
-
-#'
-#' These partitioning schemes can be used with the functions `partition_by_key()`
-#' and `partition_by_max_size()`. See Examples below.
-#'
-#' Writing a partitioned output usually requires setting `mkdir = TRUE` to
-#' automatically create the required subfolders.
+#' It is possible to export data to multiple files based on various parameters,
+#' such as the values of some variables, or such that each file has a maximum
+#' number of rows. See [partition_by()] for more details.
 #'
 #' @return The input LazyFrame.
 #' @export
@@ -124,22 +100,6 @@
 #' out_path <- withr::local_tempdir()
 #' sink_parquet(my_lf, partition_by_max_size(out_path, max_size = 5), mkdir = TRUE)
 #' fs::dir_tree(out_path) # mtcars has 32 rows so we have 7 output files
-
-# TODO: add this back when https://github.com/pola-rs/r-polars/issues/1522 is
-# solved
-#
-# # Split the LazyFrame by pre-sorted data:
-# out_path <- withr::local_tempdir()
-# my_lf |>
-#   arrange(am, cyl) |>
-#   sink_parquet(partition_parted(out_path, by = c("am", "cyl")), mkdir = TRUE)
-#
-# fs::dir_tree(out_path)
-#
-# # Careful when using partition_parted(): if the data is not presorted then
-# # the output files may be incorrect!
-# out_path <- withr::local_tempdir()
-# sink_parquet(my_lf, partition_parted(out_path, by = c("am", "cyl")), mkdir = TRUE)
 sink_parquet <- function(
   .data,
   path,
@@ -264,22 +224,6 @@ sink_parquet <- function(
 #' out_path <- withr::local_tempdir()
 #' sink_csv(my_lf, partition_by_max_size(out_path, max_size = 5), mkdir = TRUE)
 #' fs::dir_tree(out_path) # mtcars has 32 rows so we have 7 output files
-
-# TODO: add this back when https://github.com/pola-rs/r-polars/issues/1522 is
-# solved
-#
-# # Split the LazyFrame by pre-sorted data:
-# out_path <- withr::local_tempdir()
-# my_lf |>
-#   arrange(am, cyl) |>
-#   sink_csv(partition_parted(out_path, by = c("am", "cyl")), mkdir = TRUE)
-#
-# fs::dir_tree(out_path)
-#
-# # Careful when using partition_parted(): if the data is not presorted then
-# # the output files may be incorrect!
-# out_path <- withr::local_tempdir()
-# sink_csv(my_lf, partition_parted(out_path, by = c("am", "cyl")), mkdir = TRUE)
 sink_csv <- function(
   .data,
   path,
@@ -417,22 +361,6 @@ sink_csv <- function(
 #' out_path <- withr::local_tempdir()
 #' sink_ipc(my_lf, partition_by_max_size(out_path, max_size = 5), mkdir = TRUE)
 #' fs::dir_tree(out_path) # mtcars has 32 rows so we have 7 output files
-
-# TODO: add this back when https://github.com/pola-rs/r-polars/issues/1522 is
-# solved
-#
-# # Split the LazyFrame by pre-sorted data:
-# out_path <- withr::local_tempdir()
-# my_lf |>
-#   arrange(am, cyl) |>
-#   sink_ipc(partition_parted(out_path, by = c("am", "cyl")), mkdir = TRUE)
-#
-# fs::dir_tree(out_path)
-#
-# # Careful when using partition_parted(): if the data is not presorted then
-# # the output files may be incorrect!
-# out_path <- withr::local_tempdir()
-# sink_ipc(my_lf, partition_parted(out_path, by = c("am", "cyl")), mkdir = TRUE)
 sink_ipc <- function(
   .data,
   path,
@@ -520,22 +448,6 @@ sink_ipc <- function(
 #' out_path <- withr::local_tempdir()
 #' sink_ndjson(my_lf, partition_by_max_size(out_path, max_size = 5), mkdir = TRUE)
 #' fs::dir_tree(out_path) # mtcars has 32 rows so we have 7 output files
-
-# TODO: add this back when https://github.com/pola-rs/r-polars/issues/1522 is
-# solved
-#
-# # Split the LazyFrame by pre-sorted data:
-# out_path <- withr::local_tempdir()
-# my_lf |>
-#   arrange(am, cyl) |>
-#   sink_ndjson(partition_parted(out_path, by = c("am", "cyl")), mkdir = TRUE)
-#
-# fs::dir_tree(out_path)
-#
-# # Careful when using partition_parted(): if the data is not presorted then
-# # the output files may be incorrect!
-# out_path <- withr::local_tempdir()
-# sink_ndjson(my_lf, partition_parted(out_path, by = c("am", "cyl")), mkdir = TRUE)
 sink_ndjson <- function(
   .data,
   path,
@@ -572,74 +484,3 @@ sink_ndjson <- function(
     mkdir = mkdir
   )
 }
-
-#' Helper functions to export a LazyFrame as a partitioned output
-#'
-#' `r lifecycle::badge("experimental")`
-#' More details and examples in the documentation of `sink_*()` functions.
-#'
-#' @inheritParams rlang::args_dots_empty
-#' @param base_path The base path for the output files. Use the `mkdir` option
-#' of the `sink_*` methods to ensure directories in the path are created.
-#' @param by Something can be coerced to a list of Polars expressions. Used to
-#' partition by.
-#' @param include_key If `TRUE` (default), include the key columns in the output
-#' files.
-#' @param per_partition_sort_by Something can be coerced to a list of Polars
-#' expressions, or `NULL` (default). Used  to sort over within each partition.
-#' Note that this might increase the memory consumption needed for each partition.
-#' @param max_size An integer-ish value indicating the maximum number of rows in
-#' each of the generated files.
-#'
-#' @name partitioned_output
-#' @export
-partition_by_key <- function(
-  base_path,
-  ...,
-  by,
-  include_key = TRUE,
-  per_partition_sort_by = NULL
-) {
-  check_dots_empty()
-  pl$PartitionByKey(
-    base_path = base_path,
-    by = by,
-    include_key = TRUE,
-    per_partition_sort_by = NULL
-  )
-}
-
-#' @rdname partitioned_output
-#' @export
-partition_by_max_size <- function(
-  base_path,
-  ...,
-  max_size,
-  per_partition_sort_by = NULL
-) {
-  check_dots_empty()
-  pl$PartitionMaxSize(
-    base_path = base_path,
-    max_size = max_size,
-    per_partition_sort_by = NULL
-  )
-}
-
-# TODO: add this back when https://github.com/pola-rs/r-polars/issues/1522 is
-# solved
-# @export
-partition_parted <- function(
-  base_path,
-  ...,
-  by,
-  include_key = TRUE,
-  per_partition_sort_by = NULL
-) {
-  check_dots_empty()
-  pl$PartitionParted(
-    base_path = base_path,
-    by = by,
-    include_key = TRUE,
-    per_partition_sort_by = NULL
-  )
-}