| 439 | #endif |
| 440 | |
| 441 | class SparsePageSource : public SparsePageSourceImpl<SparsePage> { |
| 442 | // This is the source iterator from the user. |
| 443 | DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext> iter_; |
| 444 | DMatrixProxy* proxy_; |
| 445 | std::size_t base_row_id_{0}; |
| 446 | // Total number of batches. |
| 447 | bst_idx_t n_batches_{0}; |
| 448 | |
| 449 | void Fetch() final { |
| 450 | page_ = std::make_shared<SparsePage>(); |
| 451 | // The first round of reading, this is responsible for initialization. |
| 452 | if (!this->ReadCache()) { |
| 453 | bool type_error{false}; |
| 454 | CHECK(proxy_); |
| 455 | cpu_impl::DispatchAny( |
| 456 | proxy_, |
| 457 | [&](auto const& adapter_batch) { |
| 458 | page_->Push(adapter_batch, this->missing_, this->nthreads_); |
| 459 | }, |
| 460 | &type_error); |
| 461 | if (type_error) { |
| 462 | DevicePush(proxy_, missing_, page_.get()); |
| 463 | } |
| 464 | |
| 465 | page_->SetBaseRowId(base_row_id_); |
| 466 | base_row_id_ += page_->Size(); |
| 467 | this->n_batches_++; |
| 468 | this->WriteCache(); |
| 469 | } |
| 470 | } |
| 471 | |
| 472 | public: |
| 473 | SparsePageSource(DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext> iter, |
| 474 | DMatrixProxy* proxy, float missing, int nthreads, bst_feature_t n_features, |
| 475 | bst_idx_t n_batches, std::shared_ptr<Cache> cache) |
| 476 | : SparsePageSourceImpl(missing, nthreads, n_features, cache), |
| 477 | iter_{std::move(iter)}, |
| 478 | proxy_{proxy}, |
| 479 | n_batches_{n_batches} { |
| 480 | if (!cache_info_->written) { |
| 481 | iter_.Reset(); |
| 482 | CHECK(iter_.Next()) << "Must have at least 1 batch."; |
| 483 | } |
| 484 | this->Fetch(); |
| 485 | } |
| 486 | |
| 487 | SparsePageSource& operator++() final { |
| 488 | TryLockGuard guard{single_threaded_}; |
| 489 | count_++; |
| 490 | |
| 491 | if (cache_info_->written) { |
| 492 | at_end_ = (count_ == n_batches_); |
| 493 | } else { |
| 494 | at_end_ = !iter_.Next(); |
| 495 | } |
| 496 | CHECK_LE(count_, n_batches_); |
| 497 | |
| 498 | if (at_end_) { |