Fix matrix attributes not sliced (#4311)

dmlc · Apr 10, 2019 · 956e73f · 956e73f
1 parent 5c25755
commit 956e73f
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/R-package/R/xgb.DMatrix.R b/R-package/R/xgb.DMatrix.R
@@ -301,12 +301,17 @@ slice.xgb.DMatrix <- function(object, idxset, ...) {
 
   attr_list <- attributes(object)
   nr <- nrow(object)
-  len <- sapply(attr_list, length)
+  len <- sapply(attr_list, NROW)
   ind <- which(len == nr)
   if (length(ind) > 0) {
     nms <- names(attr_list)[ind]
     for (i in seq_along(ind)) {
-      attr(ret, nms[i]) <- attr(object, nms[i])[idxset]
+      obj_attr <- attr(object, nms[i])
+      if (NCOL(obj_attr) > 1) {
+        attr(ret, nms[i]) <- obj_attr[idxset,]
+      } else {
+        attr(ret, nms[i]) <- obj_attr[idxset]
+      }
     }
   }
   return(structure(ret, class = "xgb.DMatrix"))