speed up collector() helper (#657)

Author: simonpcouch

R/collect.R

@@ -366,11 +366,22 @@ collector <- function(x, coll_col = ".predictions") {
   } else {
     keep_cols <- coll_col
   }
-  x <- dplyr::select(x, dplyr::starts_with("id"), !!!keep_cols)
-  x <- tidyr::unnest(x, cols = c(dplyr::all_of(coll_col)))
+
+  id_cols <- colnames(x)[grepl("id", colnames(x))]
+  keep_cols <- c(coll_col, id_cols)
+  x <- x[keep_cols]
+  coll_col <- x[[coll_col]]
+
+  res <-
+    vctrs::vec_cbind(
+      vctrs::list_unchop(coll_col),
+      vctrs::vec_rep_each(x[, id_cols], times = vctrs::list_sizes(coll_col))
+    )
+
   arrange_cols <- c(".iter", ".config")
-  arrange_cols <- arrange_cols[(arrange_cols %in% names(x))]
-  arrange(x, !!!rlang::syms(arrange_cols))
+  arrange_cols <- arrange_cols[rlang::has_name(res, arrange_cols)]
+
+  res <- vctrs::vec_slice(res, vctrs::vec_order(res[arrange_cols]))
 }
 
 #' @export