From 3ef8f1f07e17ee10cc72b97666911efbcf21799a Mon Sep 17 00:00:00 2001
From: Jean-Francois Zinque <jzinque@gmail.com>
Date: Fri, 29 Mar 2019 13:17:30 +0100
Subject: [PATCH] Fix matrix attributes not sliced

---
 R-package/R/xgb.DMatrix.R | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/R-package/R/xgb.DMatrix.R b/R-package/R/xgb.DMatrix.R
index cb975cf069da..ea9cfcc479f0 100644
--- 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"))