Refactor DMatrix to return batches of different page types (#4686)

* Use explicit template parameter for specifying page type.
This commit is contained in:
Rong Ou
2019-08-03 12:10:34 -07:00
committed by Jiaming Yuan
parent e930a8e54f
commit 6edddd7966
41 changed files with 477 additions and 470 deletions

View File

@@ -28,12 +28,12 @@ TEST(SimpleDMatrix, RowAccess) {
// Loop over the batches and count the records
int64_t row_count = 0;
for (auto &batch : dmat->GetRowBatches()) {
for (auto &batch : dmat->GetBatches<xgboost::SparsePage>()) {
row_count += batch.Size();
}
EXPECT_EQ(row_count, dmat->Info().num_row_);
// Test the data read into the first row
auto &batch = *dmat->GetRowBatches().begin();
auto &batch = *dmat->GetBatches<xgboost::SparsePage>().begin();
auto first_row = batch[0];
ASSERT_EQ(first_row.size(), 3);
EXPECT_EQ(first_row[2].index, 2);
@@ -55,7 +55,7 @@ TEST(SimpleDMatrix, ColAccessWithoutBatches) {
// Loop over the batches and assert the data is as expected
int64_t num_col_batch = 0;
for (const auto &batch : dmat->GetSortedColumnBatches()) {
for (const auto &batch : dmat->GetBatches<xgboost::SortedCSCPage>()) {
num_col_batch += 1;
EXPECT_EQ(batch.Size(), dmat->Info().num_col_)
<< "Expected batch size = number of cells as #batches is 1.";