feat: Implement pl$collect_all() (#1598)

Co-authored-by: eitsupi <[email protected]>

Author: etiennebacher

DESCRIPTION

@@ -52,7 +52,7 @@ VignetteBuilder:
 Config/Needs/dev: devtools, lifecycle, readr, glue, RcppTOML, smvr
 Config/Needs/lint: fs, lintr
 Config/Needs/website: etiennebacher/altdoc, future.apply
-Config/polars/lib-version: 1.5.0-rc.2
+Config/polars/lib-version: 1.6.0-rc.1
 Config/testthat/edition: 3
 Config/testthat/parallel: true
 Config/testthat/start-first: lazyframe-frame, *-s3-base, polars_options,

NAMESPACE

@@ -252,6 +252,7 @@ export(as_polars_expr)
 export(as_polars_lf)
 export(as_polars_series)
 export(check_list_of_polars_dtype)
+export(check_list_of_polars_lf)
 export(check_polars_df)
 export(check_polars_dtype)
 export(check_polars_dtype_expr)
@@ -265,6 +266,7 @@ export(infer_polars_dtype)
 export(is_convertible_to_polars_expr)
 export(is_convertible_to_polars_series)
 export(is_list_of_polars_dtype)
+export(is_list_of_polars_lf)
 export(is_polars_df)
 export(is_polars_dtype)
 export(is_polars_dtype_expr)

NEWS.md

@@ -2,6 +2,10 @@
 
 ## polars (development version)
 
+### New features
+
+- `pl$collect_all()` to efficiently collect a list of LazyFrames (#1598).
+
 ## polars 1.5.0
 
 This is an update that corresponds to Python Polars 1.35.1.

R/000-wrappers.R

@@ -74,6 +74,11 @@ NULL
 }
 
 
+`collect_all` <- function(`lfs`, `engine`, `optflags`) {
+  .Call(savvy_collect_all__impl, `lfs`, `engine`, `optflags`)
+}
+
+
 `cols` <- function(`names`) {
   .savvy_wrap_PlRExpr(.Call(savvy_cols__impl, `names`))
 }

R/check_polars.R

@@ -108,6 +108,21 @@ is_list_of_polars_dtype <- function(x, n = NULL) {
   is_list(x, n = n) && is_list_of_polars_dtype_impl(x)
 }
 
+#' @rdname check_polars
+#' @export
+is_list_of_polars_lf <- function(x, n = NULL) {
+  is_list_of_polars_lf_impl <- function(x) {
+    for (i in seq_along(x)) {
+      if (!is_polars_lf(x[[i]])) {
+        return(FALSE)
+      }
+    }
+    TRUE
+  }
+
+  is_list(x, n = n) && is_list_of_polars_lf_impl(x)
+}
+
 #' @rdname check_polars
 #' @export
 check_polars_df <- function(
@@ -371,3 +386,32 @@ check_list_of_polars_dtype <- function(
     call = call
   )
 }
+
+#' @rdname check_polars
+#' @export
+check_list_of_polars_lf <- function(
+  x,
+  ...,
+  allow_null = FALSE,
+  arg = caller_arg(x),
+  call = caller_env()
+) {
+  if (!missing(x)) {
+    if (is_list_of_polars_lf(x)) {
+      return(invisible(NULL))
+    }
+    if (allow_null && is_null(x)) {
+      return(invisible(NULL))
+    }
+  }
+
+  stop_input_type(
+    x,
+    "a list of polars lazyframes",
+    ...,
+    allow_na = FALSE,
+    allow_null = allow_null,
+    arg = arg,
+    call = call
+  )
+}

R/functions-lazy.R

@@ -121,3 +121,54 @@ pl__arg_sort_by <- function(
     )
   })
 }
+
+#' Collect multiple LazyFrames at the same time
+#'
+#' This can run all the computation graphs in parallel or combined. Common
+#' Subplan Elimination is applied on the combined plan, meaning that diverging
+#' queries will run only once.
+#'
+#' @inheritParams rlang::args_dots_empty
+#' @inheritParams lazyframe__collect
+#' @param lazy_frames A list of LazyFrames to collect.
+#'
+#' @return A list containing all the collected DataFrames, in the same order
+#' as the input LazyFrames.
+#' @examples
+#' lf <- as_polars_lf(mtcars)$with_columns(sqrt_mpg = pl$col("mpg")$sqrt())
+#'
+#' cyl_4 <- lf$filter(pl$col("cyl") == 4)
+#' cyl_6 <- lf$filter(pl$col("cyl") == 6)
+#'
+#' # We could do `cyl_4$collect()` and `cyl_6$collect()`, but this would be
+#' # wasteful because `sqrt_mpg` would be computed twice.
+#' # `pl$collect_all()` executes only once the parts of the query that are
+#' # identical across LazyFrames.
+#' pl$collect_all(list(cyl_4, cyl_6))
+pl__collect_all <- function(
+  lazy_frames,
+  ...,
+  engine = c("auto", "in-memory", "streaming")
+) {
+  wrap({
+    check_dots_empty0(...)
+    check_list_of_polars_lf(lazy_frames)
+    engine <- arg_match0(engine, c("auto", "in-memory", "streaming"))
+
+    lfs <- lapply(lazy_frames, \(x) x$`_ldf`)
+    # TODO: add support for argument `optimizations`
+    optflags <- list(
+      comm_subexpr_elim = TRUE,
+      comm_subplan_elim = TRUE,
+      cluster_with_columns = TRUE,
+      predicate_pushdown = TRUE,
+      projection_pushdown = TRUE,
+      simplify_expression = TRUE,
+      slice_pushdown = TRUE,
+      type_coercion = TRUE
+    )
+
+    collect_all(lfs, engine = engine, optflags) |>
+      lapply(\(ptr) .savvy_wrap_PlRDataFrame(ptr) |> wrap())
+  })
+}

altdoc/mkdocs.yml

@@ -706,6 +706,7 @@ nav:
         - inner_dtype: man/datatype_expr__inner_dtype.md
         - matches: man/datatype_expr__matches.md
       Functions:
+        - collect_all: man/pl__collect_all.md
         - concat: man/pl__concat.md
         - Expr: man/polars_expr.md
         - DataFrame: man/pl__DataFrame.md