make bake() work with sparse matrices

tidymodels · Sep 10, 2024 · d5afc84 · d5afc84
1 parent 1928571
commit d5afc84
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 1 deletion.
diff --git a/NEWS.md b/NEWS.md
@@ -4,7 +4,7 @@
 
 * `recipe()`, `prep()`, and `bake()` now work with sparse tibbles. (#1364, #1366)
 
-* `recipe()` and `prep()` now work with sparse matrices. (#1364, #1368)
+* `recipe()`, `prep()`, and `bake()` now work with sparse matrices. (#1364, #1368, #1369)
 
 # recipes 1.1.0
 

diff --git a/R/recipe.R b/R/recipe.R
@@ -697,6 +697,10 @@ bake.recipe <- function(object, new_data, ..., composition = "tibble") {
     }
   }
 
+  if (is_sparse_matrix(new_data)) {
+    new_data <- sparsevctrs::coerce_to_sparse_tibble(new_data)
+  }
+
   if (!is_tibble(new_data)) {
     new_data <- as_tibble(new_data)
   }

diff --git a/tests/testthat/test-sparsevctrs.R b/tests/testthat/test-sparsevctrs.R
@@ -133,3 +133,28 @@ test_that("prep() accepts sparse matrices", {
     is_sparse_tibble(rec$template)
   )
 })
+
+test_that("bake() accepts sparse matrices", {
+  skip_if_not_installed("modeldata")
+
+  hotel_data <- sparse_hotel_rates()
+
+  rec_spec <- recipe(avg_price_per_room ~ ., data = hotel_data) %>%
+    prep()
+
+  expect_no_condition(
+    res <- bake(rec_spec, new_data = NULL)
+  )
+
+  expect_true(
+    is_sparse_tibble(res)
+  )
+
+  expect_no_error(
+    res <- bake(rec_spec, new_data = hotel_data)
+  )
+
+  expect_true(
+    is_sparse_tibble(res)
+  )
+})