Function TestColumnSplitPartitioner

tests/cpp/tree/test_approx.cc:135–182

namespace {
void TestColumnSplitPartitioner(size_t n_samples, size_t base_rowid, std::shared_ptr<DMatrix> Xy,
                                std::vector<float>* hess, float min_value, float mid_value,
                                CommonRowPartitioner const& expected_mid_partitioner) {
  auto dmat =
      std::unique_ptr<DMatrix>{Xy->SliceCol(collective::GetWorldSize(), collective::GetRank())};
  std::vector<CPUExpandEntry> candidates{{0, 0}};
  candidates.front().split.loss_chg = 0.4;

  Context ctx;
  ctx.InitAllowUnknown(Args{});
  for (auto const& page : dmat->GetBatches<GHistIndexMatrix>(&ctx, {64, *hess, true})) {
    {
      RegTree tree;
      CommonRowPartitioner partitioner{&ctx, n_samples, base_rowid, true};
      GetSplit(&tree, min_value, &candidates);
      partitioner.UpdatePosition(&ctx, page, candidates, tree.HostScView());
      ASSERT_EQ(partitioner.Size(), 3);
      ASSERT_EQ(partitioner[1].Size(), 0);
      ASSERT_EQ(partitioner[2].Size(), n_samples);
    }
    {
      CommonRowPartitioner partitioner{&ctx, n_samples, base_rowid, true};
      RegTree tree;
      GetSplit(&tree, mid_value, &candidates);
      partitioner.UpdatePosition(&ctx, page, candidates, tree.HostScView());
      {
        auto left_nidx = tree[RegTree::kRoot].LeftChild();
        auto const& elem = partitioner[left_nidx];
        ASSERT_LT(elem.Size(), n_samples);
        ASSERT_GT(elem.Size(), 1);
        auto const& expected_elem = expected_mid_partitioner[left_nidx];
        ASSERT_EQ(elem.Size(), expected_elem.Size());
        for (auto it = elem.begin(), eit = expected_elem.begin(); it != elem.end(); ++it, ++eit) {
          ASSERT_EQ(*it, *eit);
        }
      }
      {
        auto right_nidx = tree[RegTree::kRoot].RightChild();
        auto const& elem = partitioner[right_nidx];
        auto const& expected_elem = expected_mid_partitioner[right_nidx];
        ASSERT_EQ(elem.Size(), expected_elem.Size());
        for (auto it = elem.begin(), eit = expected_elem.begin(); it != elem.end(); ++it, ++eit) {
          ASSERT_EQ(*it, *eit);
        }
      }
    }
  }
}
}  // anonymous namespace

TEST(Approx, PartitionerColumnSplit) {