fix dmatrix unit test

Author: vaaaaanquish

src/dmatrix.rs

@@ -7,7 +7,8 @@ use xgboost_sys;
 
 use super::{XGBResult, XGBError};
 
-static KEY_ROOT_INDEX: &'static str = "root_index";
+static KEY_GROUP_PTR: &'static str = "group_ptr";
+static KEY_GROUP: &'static str = "group";
 static KEY_LABEL: &'static str = "label";
 static KEY_WEIGHT: &'static str = "weight";
 static KEY_BASE_MARGIN: &'static str = "base_margin";
@@ -230,20 +231,6 @@ impl DMatrix {
         Ok(DMatrix::new(out_handle)?)
     }
 
-    /// Gets the specified root index of each instance, can be used for multi task setting.
-    ///
-    /// See the XGBoost documentation for more information.
-    pub fn get_root_index(&self) -> XGBResult<&[u32]> {
-        self.get_uint_info(KEY_ROOT_INDEX)
-    }
-
-    /// Sets the specified root index of each instance, can be used for multi task setting.
-    ///
-    /// See the XGBoost documentation for more information.
-    pub fn set_root_index(&mut self, array: &[u32]) -> XGBResult<()> {
-        self.set_uint_info(KEY_ROOT_INDEX, array)
-    }
-
     /// Get ground truth labels for each row of this matrix.
     pub fn get_labels(&self) -> XGBResult<&[f32]> {
         self.get_float_info(KEY_LABEL)
@@ -282,9 +269,20 @@ impl DMatrix {
     ///
     /// See the XGBoost documentation for more information.
     pub fn set_group(&mut self, group: &[u32]) -> XGBResult<()> {
-        xgb_call!(xgboost_sys::XGDMatrixSetGroup(self.handle, group.as_ptr(), group.len() as u64))
+        // same as xgb_call!(xgboost_sys::XGDMatrixSetGroup(self.handle, group.as_ptr(), group.len() as u64))
+        self.set_uint_info(KEY_GROUP, group)
+    }
+
+    /// Get the index for the beginning and end of a group.
+    ///
+    /// Needed when the learning task is ranking.
+    ///
+    /// See the XGBoost documentation for more information.
+    pub fn get_group(&self) -> XGBResult<&[u32]> {
+        self.get_uint_info(KEY_GROUP_PTR)
     }
 
+
     fn get_float_info(&self, field: &str) -> XGBResult<&[f32]> {
         let field = ffi::CString::new(field).unwrap();
         let mut out_len = 0;
@@ -313,7 +311,6 @@ impl DMatrix {
                                                     field.as_ptr(),
                                                     &mut out_len,
                                                     &mut out_dptr))?;
-
         Ok(unsafe { slice::from_raw_parts(out_dptr as *mut c_uint, out_len as usize) })
     }
 
@@ -370,16 +367,6 @@ mod tests {
         // TODO: check contents as well, if possible
     }
 
-    #[test]
-    fn get_set_root_index() {
-        let mut dmat = read_train_matrix().unwrap();
-        assert_eq!(dmat.get_root_index().unwrap(), &[]);
-
-        let root_index = [3, 22, 1];
-        assert!(dmat.set_root_index(&root_index).is_ok());
-        assert_eq!(dmat.get_root_index().unwrap(), &[3, 22, 1]);
-    }
-
     #[test]
     fn get_set_labels() {
         let mut dmat = read_train_matrix().unwrap();
@@ -395,7 +382,7 @@ mod tests {
         let mut dmat = read_train_matrix().unwrap();
         assert_eq!(dmat.get_weights().unwrap(), &[]);
 
-        let weight = [1.0, 10.0, -123.456789, 44.9555];
+        let weight = [1.0, 10.0, 44.9555];
         assert!(dmat.set_weights(&weight).is_ok());
         assert_eq!(dmat.get_weights().unwrap(), weight);
     }
@@ -411,11 +398,13 @@ mod tests {
     }
 
     #[test]
-    fn set_group() {
+    fn get_set_group() {
         let mut dmat = read_train_matrix().unwrap();
+        assert_eq!(dmat.get_group().unwrap(), &[]);
 
-        let group = [1, 2, 3];
+        let group = [1];
         assert!(dmat.set_group(&group).is_ok());
+        assert_eq!(dmat.get_group().unwrap(), &[0, 1]);
     }
 
     #[test]
@@ -426,7 +415,7 @@ mod tests {
 
         let dmat = DMatrix::from_csr(&indptr, &indices, &data, None).unwrap();
         assert_eq!(dmat.num_rows(), 4);
-        assert_eq!(dmat.num_cols(), 3);
+        assert_eq!(dmat.num_cols(), 0);  // https://github.com/dmlc/xgboost/pull/7265
 
         let dmat = DMatrix::from_csr(&indptr, &indices, &data, Some(10)).unwrap();
         assert_eq!(dmat.num_rows(), 4);
@@ -477,7 +466,7 @@ mod tests {
         assert_eq!(dmat.slice(&[1]).unwrap().shape(), (1, 2));
         assert_eq!(dmat.slice(&[0, 1]).unwrap().shape(), (2, 2));
         assert_eq!(dmat.slice(&[3, 2, 1]).unwrap().shape(), (3, 2));
-        assert!(dmat.slice(&[10, 11, 12]).is_err());
+        assert_eq!(dmat.slice(&[10, 11, 12]).unwrap().shape(), (0, 0));
     }
 
     #[test]