diff --git a/R-package/R/xgb.model.dt.tree.R b/R-package/R/xgb.model.dt.tree.R index 598132161511..5411c35d2394 100644 --- a/R-package/R/xgb.model.dt.tree.R +++ b/R-package/R/xgb.model.dt.tree.R @@ -134,7 +134,10 @@ xgb.model.dt.tree <- function(feature_names = NULL, model = NULL, text = NULL, ] # assign feature_names when available - if (!is.null(feature_names) && !is.na(td$Feature)) { + is_stump <- function() { + return(length(td$Feature) == 1 && is.na(td$Feature)) + } + if (!is.null(feature_names) && !is_stump()) { if (length(feature_names) <= max(as.numeric(td$Feature), na.rm = TRUE)) stop("feature_names has less elements than there are features used in the model") td[isLeaf == FALSE, Feature := feature_names[as.numeric(Feature) + 1]] diff --git a/R-package/tests/testthat/test_helpers.R b/R-package/tests/testthat/test_helpers.R index c426c8377468..fdd0ce02b5a6 100644 --- a/R-package/tests/testthat/test_helpers.R +++ b/R-package/tests/testthat/test_helpers.R @@ -340,6 +340,16 @@ test_that("xgb.importance works with and without feature names", { imp } expect_equal(importance_from_dump(), importance, tolerance = 1e-6) + + ## decision stump + m <- xgboost::xgboost( + data = as.matrix(data.frame(x = c(0, 1))), + label = c(1, 2), + nrounds = 1 + ) + df <- xgb.model.dt.tree(model = m) + expect_equal(df$Feature, "Leaf") + expect_equal(df$Cover, 2) }) test_that("xgb.importance works with GLM model", {