Refactor SparsePageSource, delete cache files after use

src/data/ellpack_page_source.cu

@@ -45,7 +45,7 @@ class EllpackPageSourceImpl : public DataSource<EllpackPage> {
   dh::BulkAllocator ba_;
   /*! \brief The EllpackInfo, with the underlying GPU memory shared by all pages. */
   EllpackInfo ellpack_info_;
-  std::unique_ptr<SparsePageSource<EllpackPage>> source_;
+  std::unique_ptr<ExternalMemoryPrefetcher<EllpackPage>> source_;
   std::string cache_info_;
 };
 
@@ -98,11 +98,13 @@ EllpackPageSourceImpl::EllpackPageSourceImpl(DMatrix* dmat,
   WriteEllpackPages(dmat, cache_info);
   monitor_.StopCuda("WriteEllpackPages");
 
-  source_.reset(new SparsePageSource<EllpackPage>(cache_info_, kPageType_));
+  source_.reset(new ExternalMemoryPrefetcher<EllpackPage>(
+      ParseCacheInfo(cache_info_, kPageType_)));
 }
 
 void EllpackPageSourceImpl::BeforeFirst() {
-  source_.reset(new SparsePageSource<EllpackPage>(cache_info_, kPageType_));
+  source_.reset(new ExternalMemoryPrefetcher<EllpackPage>(
+      ParseCacheInfo(cache_info_, kPageType_)));
   source_->BeforeFirst();
 }
 

src/data/sparse_page_dmatrix.cc

@@ -23,58 +23,24 @@ const MetaInfo& SparsePageDMatrix::Info() const {
   return row_source_->info;
 }
 
-template<typename S, typename T>
-class SparseBatchIteratorImpl : public BatchIteratorImpl<T> {
- public:
-  explicit SparseBatchIteratorImpl(S* source) : source_(source) {
-    CHECK(source_ != nullptr);
-  }
-  T& operator*() override { return source_->Value(); }
-  const T& operator*() const override { return source_->Value(); }
-  void operator++() override { at_end_ = !source_->Next(); }
-  bool AtEnd() const override { return at_end_; }
-
- private:
-  S* source_{nullptr};
-  bool at_end_{ false };
-};
-
 BatchSet<SparsePage> SparsePageDMatrix::GetRowBatches() {
-  auto cast = dynamic_cast<SparsePageSource<SparsePage>*>(row_source_.get());
-  CHECK(cast);
-  cast->BeforeFirst();
-  cast->Next();
-  auto begin_iter = BatchIterator<SparsePage>(
-      new SparseBatchIteratorImpl<SparsePageSource<SparsePage>, SparsePage>(cast));
-  return BatchSet<SparsePage>(begin_iter);
+  return row_source_->GetBatchSet();
 }
 
 BatchSet<CSCPage> SparsePageDMatrix::GetColumnBatches() {
   // Lazily instantiate
   if (!column_source_) {
-    SparsePageSource<SparsePage>::CreateColumnPage(this, cache_info_, false);
-    column_source_.reset(new SparsePageSource<CSCPage>(cache_info_, ".col.page"));
+    column_source_.reset(new CSCPageSource(this, cache_info_));
   }
-  column_source_->BeforeFirst();
-  column_source_->Next();
-  auto begin_iter = BatchIterator<CSCPage>(
-      new SparseBatchIteratorImpl<SparsePageSource<CSCPage>, CSCPage>(column_source_.get()));
-  return BatchSet<CSCPage>(begin_iter);
+  return column_source_->GetBatchSet();
 }
 
 BatchSet<SortedCSCPage> SparsePageDMatrix::GetSortedColumnBatches() {
   // Lazily instantiate
   if (!sorted_column_source_) {
-    SparsePageSource<SparsePage>::CreateColumnPage(this, cache_info_, true);
-    sorted_column_source_.reset(
-        new SparsePageSource<SortedCSCPage>(cache_info_, ".sorted.col.page"));
+    sorted_column_source_.reset(new SortedCSCPageSource(this, cache_info_));
   }
-  sorted_column_source_->BeforeFirst();
-  sorted_column_source_->Next();
-  auto begin_iter = BatchIterator<SortedCSCPage>(
-      new SparseBatchIteratorImpl<SparsePageSource<SortedCSCPage>, SortedCSCPage>(
-          sorted_column_source_.get()));
-  return BatchSet<SortedCSCPage>(begin_iter);
+  return sorted_column_source_->GetBatchSet();
 }
 
 BatchSet<EllpackPage> SparsePageDMatrix::GetEllpackBatches(const BatchParam& param) {

src/data/sparse_page_dmatrix.h

@@ -22,24 +22,15 @@ namespace data {
 // Used for external memory.
 class SparsePageDMatrix : public DMatrix {
  public:
-  explicit SparsePageDMatrix(std::unique_ptr<DataSource<SparsePage>>&& source,
-                             std::string cache_info)
-      : row_source_(std::move(source)), cache_info_(std::move(cache_info)) {}
-
   template <typename AdapterT>
   explicit SparsePageDMatrix(AdapterT* adapter, float missing, int nthread,
                              const std::string& cache_prefix,
                              size_t page_size = kPageSize)
       : cache_info_(std::move(cache_prefix)) {
-    if (!data::SparsePageSource<SparsePage>::CacheExist(cache_prefix,
-                                                        ".row.page")) {
-      data::SparsePageSource<SparsePage>::CreateRowPage(
-          adapter, missing, nthread, cache_prefix, page_size);
-    }
-    row_source_.reset(
-        new data::SparsePageSource<SparsePage>(cache_prefix, ".row.page"));
+    row_source_.reset(new data::SparsePageSource(adapter, missing, nthread,
+                                                 cache_prefix, page_size));
   }
-    // Set number of threads but keep old value so we can reset it after
+  // Set number of threads but keep old value so we can reset it after
   ~SparsePageDMatrix() override = default;
 
   MetaInfo& Info() override;
@@ -57,9 +48,9 @@ class SparsePageDMatrix : public DMatrix {
   BatchSet<EllpackPage> GetEllpackBatches(const BatchParam& param) override;
 
   // source data pointers.
-  std::unique_ptr<DataSource<SparsePage>> row_source_;
-  std::unique_ptr<SparsePageSource<CSCPage>> column_source_;
-  std::unique_ptr<SparsePageSource<SortedCSCPage>> sorted_column_source_;
+  std::unique_ptr<SparsePageSource> row_source_;
+  std::unique_ptr<CSCPageSource> column_source_;
+  std::unique_ptr<SortedCSCPageSource> sorted_column_source_;
   std::unique_ptr<EllpackPageSource> ellpack_source_;
   // saved batch param
   BatchParam batch_param_;