pass along nfeatures and feature_names to xgboost bundle (#67)

EmilHvitfeldt · juliasilge · web-flow · commit 51fbfb425130 · 2024-11-08T10:54:45.000-07:00
* pass along nfeatures and feature_names to xgboost bundle

* add tests

* Update NEWS

---------

Co-authored-by: Julia Silge &lt;julia.silge@gmail.com&gt;
diff --git a/NEWS.md b/NEWS.md
@@ -3,6 +3,8 @@
 * Added bundle method for objects from `dbarts::bart()` and, by extension,
   `parsnip::bart(engine = "dbarts")` (#64).
 
+* Bundling xgboost objects now takes extra steps to preserve `nfeatures` and `feature_names` (#67).
+
 # bundle 0.1.1
 
 * Fixed bundling of recipes steps situated inside of workflows.
diff --git a/R/bundle_xgboost.R b/R/bundle_xgboost.R
@@ -51,6 +51,9 @@ bundle.xgb.Booster <- function(x, ...) {
         num_class = !!x$params$num_class
       )
 
+      res$nfeatures <- !!x$nfeatures
+      res$feature_names <- !!x$feature_names
+
       res
     }),
     desc_class = class(x)[1]
diff --git a/tests/testthat/test_bundle_xgboost.R b/tests/testthat/test_bundle_xgboost.R
@@ -96,4 +96,8 @@ test_that("bundling + unbundling xgboost fits", {
   # compare predictions
   expect_equal(mod_preds, mod_unbundled_preds)
   expect_equal(mod_preds, mod_butchered_unbundled_preds)
+
+  # verify nfeatures and feature_names are kept
+  expect_identical(unbundle(mod_bundle)$nfeatures, mod_fit$nfeatures)
+  expect_identical(unbundle(mod_bundle)$feature_names, mod_fit$feature_names)
 })

Original file line number	Diff line number	Diff line change
`@@ -51,6 +51,9 @@ bundle.xgb.Booster <- function(x, ...) {`
`51`	`51`	`num_class = !!x$params$num_class`
`52`	`52`	`)`
`53`	`53`
	`54`	`+ res$nfeatures <- !!x$nfeatures`
	`55`	`+ res$feature_names <- !!x$feature_names`
	`56`	`+`
`54`	`57`	`res`
`55`	`58`	`}),`
`56`	`59`	`desc_class = class(x)[1]`