Fix ctor.

dmlc · Jul 4, 2021 · ff82526 · ff82526
1 parent 72e5b35
commit ff82526
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/src/data/sparse_page_dmatrix.cc b/src/data/sparse_page_dmatrix.cc
@@ -41,7 +41,7 @@ SparsePageDMatrix::SparsePageDMatrix(DataIterHandle iter_handle, DMatrixHandle p
         proxy, [](auto const &value) { return value.NumCols(); });
   };
 
-  for (auto const &page : this->GetRowBatches()) {
+  for (auto const &page : this->GetRowBatchesImpl()) {
     this->info_.Extend(std::move(proxy->Info()), false, false);
     n_features = std::max(n_features, num_cols());
     n_samples += num_rows();
@@ -77,12 +77,16 @@ void SparsePageDMatrix::InitializeSparsePage() {
       this->n_batches_, cache_info_.at(id));
 }
 
-BatchSet<SparsePage> SparsePageDMatrix::GetRowBatches() {
+BatchSet<SparsePage> SparsePageDMatrix::GetRowBatchesImpl() {
   this->InitializeSparsePage();
   auto begin_iter = BatchIterator<SparsePage>(sparse_page_source_);
   return BatchSet<SparsePage>(BatchIterator<SparsePage>(begin_iter));
 }
 
+BatchSet<SparsePage> SparsePageDMatrix::GetRowBatches() {
+  return this->GetRowBatchesImpl();
+}
+
 BatchSet<CSCPage> SparsePageDMatrix::GetColumnBatches() {
   auto id = MakeCache(this, ".col.page", cache_prefix_, &cache_info_);
   CHECK_NE(this->Info().num_col_, 0);

diff --git a/src/data/sparse_page_dmatrix.h b/src/data/sparse_page_dmatrix.h
@@ -39,6 +39,8 @@ class SparsePageDMatrix : public DMatrix {
   size_t n_batches_ {0};
 
   void InitializeSparsePage();
+  // Non virtual version that can be used in constructor
+  BatchSet<SparsePage> GetRowBatchesImpl();
 
  public:
   explicit SparsePageDMatrix(DataIterHandle iter, DMatrixHandle proxy,