Implement slice via adapters (#5198)

This commit is contained in:
Rory Mitchell
2020-01-14 12:55:41 +13:00
committed by GitHub
parent f100b8d878
commit a73e25e15f
7 changed files with 190 additions and 45 deletions

View File

@@ -61,3 +61,31 @@ TEST(adapter, CSCAdapterColsMoreThanRows) {
EXPECT_EQ(inst[3].fvalue, 8);
EXPECT_EQ(inst[3].index, 3);
}
TEST(c_api, DMatrixSliceAdapterFromSimpleDMatrix) {
auto pp_dmat = CreateDMatrix(6, 2, 1.0);
auto p_dmat = *pp_dmat;
std::vector<int> ridx_set = {1, 3, 5};
data::DMatrixSliceAdapter adapter(p_dmat.get(),
{ridx_set.data(), ridx_set.size()});
EXPECT_EQ(adapter.NumRows(), ridx_set.size());
adapter.BeforeFirst();
for (auto &batch : p_dmat->GetBatches<SparsePage>()) {
adapter.Next();
auto &adapter_batch = adapter.Value();
for (auto i = 0ull; i < adapter_batch.Size(); i++) {
auto inst = batch[ridx_set[i]];
auto line = adapter_batch.GetLine(i);
ASSERT_EQ(inst.size(), line.Size());
for (auto j = 0ull; j < line.Size(); j++) {
EXPECT_EQ(inst[j].fvalue, line.GetElement(j).value);
EXPECT_EQ(inst[j].index, line.GetElement(j).column_idx);
EXPECT_EQ(i, line.GetElement(j).row_idx);
}
}
}
delete pp_dmat;
}