Merge pull request #1168 from tidymodels/fix1166

save x column names from fit_xy()

Author: EmilHvitfeldt

NEWS.md

@@ -14,6 +14,7 @@
 
 * Ensure that `knit_engine_docs()` has the required packages installed (#1156).
 
+* Fixed bug where some models fit using `fit_xy()` couldn't predict (#1166).
 
 # parsnip 1.2.1
 

R/fit_helpers.R

@@ -115,7 +115,7 @@ xy_xy <- function(object,
   } else {
     y_name <- colnames(env$y)
   }
-  res$preproc <- list(y_var = y_name)
+  res$preproc <- list(y_var = y_name, x_names = colnames(env$x))
   res$elapsed <- list(elapsed = elapsed, print = control$verbosity > 1L)
   res
 }

R/predict.R

@@ -440,11 +440,17 @@ prepare_data <- function(object, new_data) {
   preproc_names <- names(object$preproc)
   translate_from_formula_to_xy <- any(preproc_names == "terms", na.rm = TRUE)
   translate_from_xy_to_formula <- any(preproc_names == "x_var", na.rm = TRUE)
+  # For backwards compatibility, only do this if `y_var` is missing and
+  # `x_names` is present
+  translate_from_xy_to_xy <- any(preproc_names == "x_names", na.rm = TRUE) &&
+    identical(object$preproc$y_var, character(0))
 
   if (translate_from_formula_to_xy) {
     new_data <- .convert_form_to_xy_new(object$preproc, new_data)$x
   } else if (translate_from_xy_to_formula) {
     new_data <- .convert_xy_to_form_new(object$preproc, new_data)
+  } else if (translate_from_xy_to_xy) {
+    new_data <- new_data[, object$preproc$x_names]
   }
 
   encodings <- get_encoding(class(object$spec)[1])

tests/testthat/test-predict_formats.R

@@ -106,4 +106,20 @@ test_that('non-factor classification', {
   )
 })
 
+test_that("predict() works for model fit with fit_xy() (#1166)", {
+  skip_if_not_installed("xgboost")
 
+  spec <- boost_tree() %>%
+    set_mode("regression") %>%
+    set_engine("xgboost")
+
+  tree_fit <- fit(spec, mpg ~ ., data = mtcars)
+
+  exp <- predict(tree_fit, mtcars)
+
+  tree_fit <- fit_xy(spec, x = mtcars[, -1], y = mtcars[, 1])
+
+  res <- predict(tree_fit, mtcars)
+  
+  expect_identical(exp, res)
+})