Improve mdl_df_apply for extra key variables of mable

tidyverts · Sep 13, 2024 · 445d65f · 445d65f
1 parent b61d4bf
commit 445d65f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/R/utils.R b/R/utils.R
@@ -325,8 +325,9 @@ dist_types <- function(dist) {
 
 mdl_df_apply <- function(x, f, ...) {
   mbl_vars <- mable_vars(x)
+  kv <- key_vars(x)
   x <- mutate(as_tibble(x), 
               dplyr::across(all_of(mbl_vars), function(x) lapply(x, f, ...)))
   x <- pivot_longer(x, all_of(mbl_vars), names_to = ".model", values_to = "__results__")
-  unnest_tsbl(x, "__results__", parent_key = ".model")
+  unnest_tsbl(x, "__results__", parent_key = c(kv, ".model"))
 }