Fix accumulate rows.

dmlc · Jul 1, 2021 · 2617120 · 2617120
1 parent f18ecd4
commit 2617120
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 7 deletions.
diff --git a/include/xgboost/data.h b/include/xgboost/data.h
@@ -170,7 +170,7 @@ class MetaInfo {
    *
    * \param that The other MetaInfo object.
    */
-  void Extend(MetaInfo const& that, bool check_column);
+  void Extend(MetaInfo const& that, bool accumulate_rows, bool check_column);
 
  private:
   /*! \brief argsort of labels */

diff --git a/src/data/data.cc b/src/data/data.cc
@@ -499,8 +499,10 @@ void MetaInfo::GetFeatureInfo(const char *field,
   }
 }
 
-void MetaInfo::Extend(MetaInfo const& that, bool check_column) {
-  this->num_row_ += that.num_row_;
+void MetaInfo::Extend(MetaInfo const& that, bool accumulate_rows, bool check_column) {
+  if (accumulate_rows) {
+    this->num_row_ += that.num_row_;
+  }
   if (this->num_col_ != 0) {
     if (check_column) {
       CHECK_EQ(this->num_col_, that.num_col_)

diff --git a/src/data/iterative_device_dmatrix.cu b/src/data/iterative_device_dmatrix.cu
@@ -144,7 +144,7 @@ void IterativeDeviceDMatrix::Initialize(DataIterHandle iter_handle, float missin
     proxy->Info().num_row_ = num_rows();
     proxy->Info().num_col_ = cols;
     if (batches != 1) {
-      this->info_.Extend(std::move(proxy->Info()), false);
+      this->info_.Extend(std::move(proxy->Info()), false, true);
     }
     n_batches_for_verification++;
   }

diff --git a/src/data/sparse_page_dmatrix.cc b/src/data/sparse_page_dmatrix.cc
@@ -36,7 +36,7 @@ void SparsePageDMatrix::InitializeExternalMemory() {
   };
 
   for (auto const &page : this->GetRowBatches()) {
-    this->info_.Extend(std::move(proxy->Info()), false);
+    this->info_.Extend(std::move(proxy->Info()), false, false);
     n_features = std::max(n_features, num_cols());
     n_samples += num_rows();
     nnz += page.data.Size();

diff --git a/tests/cpp/data/test_metainfo.cc b/tests/cpp/data/test_metainfo.cc
@@ -260,7 +260,7 @@ TEST(MetaInfo, HostExtend) {
   lhs.SetInfo("group", groups.data(), xgboost::DataType::kUInt32, groups.size());
   rhs.SetInfo("group", groups.data(), xgboost::DataType::kUInt32, groups.size());
 
-  lhs.Extend(rhs, true);
+  lhs.Extend(rhs, true, true);
   ASSERT_EQ(lhs.num_row_, kRows * 2);
   ASSERT_TRUE(lhs.labels_.HostCanRead());
   ASSERT_TRUE(rhs.labels_.HostCanRead());

diff --git a/tests/cpp/data/test_metainfo.cu b/tests/cpp/data/test_metainfo.cu
@@ -141,7 +141,7 @@ TEST(MetaInfo, DeviceExtend) {
   lhs.num_row_ = kRows;
   rhs.num_row_ = kRows;
 
-  lhs.Extend(rhs, true);
+  lhs.Extend(rhs, true, true);
   ASSERT_EQ(lhs.num_row_, kRows * 2);
   ASSERT_FALSE(lhs.labels_.HostCanRead());