More tests.

Author: trivialfis

R-package/tests/testthat/test_dmatrix.R

@@ -59,7 +59,7 @@ test_that("xgb.DMatrix: getinfo & setinfo", {
   expect_error(setinfo(dtest, 'group', test_label))
 
   # providing character values will give a warning
-  expect_warning( setinfo(dtest, 'weight', rep('a', nrow(test_data))) )
+  expect_warning(setinfo(dtest, 'weight', rep('a', nrow(test_data))))
 
   # any other label should error
   expect_error(setinfo(dtest, 'asdf', test_label))

src/data/data.cc

@@ -223,6 +223,8 @@ std::vector<T> Gather(const std::vector<T> &in, common::Span<int const> ridxs, s
 
 MetaInfo MetaInfo::Slice(common::Span<int32_t const> ridxs) const {
   MetaInfo out;
+  out.num_row_ = ridxs.size();
+  out.num_col_ = this->num_col_;
   // Groups is maintained by a higher level Python function.  We should aim at deprecating
   // the slice function.
   out.labels_.HostVector() = Gather(this->labels_.HostVector(), ridxs);

src/data/simple_dmatrix.cc

@@ -31,8 +31,9 @@ DMatrix* SimpleDMatrix::Slice(common::Span<int32_t const> ridxs) {
       std::copy(inst.begin(), inst.end(), std::back_inserter(h_data));
       h_offset.emplace_back(rptr);
     }
+    out->Info() = this->Info().Slice(ridxs);
+    out->Info().num_nonzero_ = h_offset.back();
   }
-  out->Info() = this->Info().Slice(ridxs);
   return out;
 }
 

tests/cpp/data/test_simple_dmatrix.cc

@@ -222,6 +222,10 @@ TEST(SimpleDMatrix, Slice) {
   size_t constexpr kCols {8};
   size_t constexpr kClasses {3};
   auto p_m = RandomDataGenerator{kRows, kCols, 0}.GenerateDMatrix(true);
+  auto& weights = p_m->Info().weights_.HostVector();
+  weights.resize(kRows);
+  std::iota(weights.begin(), weights.end(), 0.0f);
+
   auto& lower = p_m->Info().labels_lower_bound_.HostVector();
   auto& upper = p_m->Info().labels_upper_bound_.HostVector();
   lower.resize(kRows);
@@ -256,6 +260,8 @@ TEST(SimpleDMatrix, Slice) {
                   out->Info().labels_lower_bound_.HostVector().at(i));
         ASSERT_EQ(p_m->Info().labels_upper_bound_.HostVector().at(ridx),
                   out->Info().labels_upper_bound_.HostVector().at(i));
+        ASSERT_EQ(p_m->Info().weights_.HostVector().at(ridx),
+                  out->Info().weights_.HostVector().at(i));
 
         auto& out_margin = out->Info().base_margin_.HostVector();
         for (size_t j = 0; j < kClasses; ++j) {
@@ -265,6 +271,10 @@ TEST(SimpleDMatrix, Slice) {
       }
     }
   }
+
+  ASSERT_EQ(out->Info().num_col_, out->Info().num_col_);
+  ASSERT_EQ(out->Info().num_row_, ridxs.size());
+  ASSERT_EQ(out->Info().num_nonzero_, ridxs.size() * kCols);  // dense
 }
 
 TEST(SimpleDMatrix, SaveLoadBinary) {

tests/python/test_dmatrix.py

@@ -71,7 +71,34 @@ def test_np_view(self):
         assert (from_view.shape == from_array.shape)
         assert (from_view == from_array).all()
 
-    def test_feature_names(self):
+    def test_slice(self):
+        X = rng.randn(100, 100)
+        y = rng.randint(low=0, high=3, size=100)
+        d = xgb.DMatrix(X, y)
+        eval_res_0 = {}
+        booster = xgb.train(
+            {'num_class': 3, 'objective': 'multi:softprob'}, d,
+            num_boost_round=2, evals=[(d, 'd')], evals_result=eval_res_0)
+
+        predt = booster.predict(d)
+        predt = predt.reshape(100 * 3, 1)
+        d.set_base_margin(predt)
+
+        ridxs = [1, 2, 3, 4, 5, 6]
+        d = d.slice(ridxs)
+        sliced_margin = d.get_float_info('base_margin')
+        assert sliced_margin.shape[0] == len(ridxs) * 3
+
+        eval_res_1 = {}
+        xgb.train({'num_class': 3, 'objective': 'multi:softprob'}, d,
+                  num_boost_round=2, evals=[(d, 'd')], evals_result=eval_res_1)
+
+        eval_res_0 = eval_res_0['d']['merror']
+        eval_res_1 = eval_res_1['d']['merror']
+        for i in range(len(eval_res_0)):
+            assert abs(eval_res_0[i] - eval_res_1[i]) < 0.02
+
+    def test_feature_names_slice(self):
         data = np.random.randn(5, 5)
 
         # different length