#include "column_matrix.h"
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <type_traits>
#include <vector>
#include "../data/gradient_index.h"
#include "io.h"
#include "xgboost/base.h"
#include "xgboost/span.h"
namespace xgboost::common {
void ColumnMatrix::InitStorage(GHistIndexMatrix const& gmat, double sparse_threshold) {
auto const nfeature = gmat.Features();
const size_t nrow = gmat.Size();
type_ = common::MakeFixedVecWithMalloc(nfeature, ColumnType{});
uint32_t max_val = std::numeric_limits<uint32_t>::max();
for (bst_feature_t fid = 0; fid < nfeature; ++fid) {
CHECK_LE(gmat.cut.Ptrs()[fid + 1] - gmat.cut.Ptrs()[fid], max_val);
}
bool all_dense_column = true;
std::vector<size_t> feature_counts(nfeature, 0);
gmat.GetFeatureCounts(feature_counts.data());
for (bst_feature_t fid = 0; fid < nfeature; ++fid) {
if (static_cast<double>(feature_counts[fid]) < sparse_threshold * nrow) {
type_[fid] = kSparseColumn;
all_dense_column = false;
} else {
type_[fid] = kDenseColumn;
}
}
feature_offsets_ = common::MakeFixedVecWithMalloc(nfeature + 1, std::size_t{0});
size_t accum_index = 0;
feature_offsets_[0] = accum_index;
for (bst_feature_t fid = 1; fid < nfeature + 1; ++fid) {
if (type_[fid - 1] == kDenseColumn) {
accum_index += static_cast<size_t>(nrow);
} else {
accum_index += feature_counts[fid - 1];
}
feature_offsets_[fid] = accum_index;
}
SetTypeSize(gmat.MaxNumBinPerFeat());
auto storage_size =
feature_offsets_.back() * static_cast<std::underlying_type_t<BinTypeSize>>(bins_type_size_);
index_ = common::MakeFixedVecWithMalloc(storage_size, std::uint8_t{0});
if (!all_dense_column) {
row_ind_ = common::MakeFixedVecWithMalloc(feature_offsets_[nfeature], std::size_t{0});
}
index_base_ = const_cast<uint32_t*>(gmat.cut.Ptrs().data());
any_missing_ = !gmat.IsDense();
missing_ = MissingIndicator{0, false};
}
bool ColumnMatrix::Read(AlignedResourceReadStream* fi, uint32_t const* index_base) {
if (!common::ReadVec(fi, &index_)) {
return false;
}
if (!common::ReadVec(fi, &type_)) {
return false;
}
if (!common::ReadVec(fi, &row_ind_)) {
return false;
}
if (!common::ReadVec(fi, &feature_offsets_)) {
return false;
}
if (!common::ReadVec(fi, &missing_.storage)) {
return false;
}
missing_.InitView();
index_base_ = index_base;
if (!fi->Read(&bins_type_size_)) {
return false;
}
if (!fi->Read(&any_missing_)) {
return false;
}
return true;
}
std::size_t ColumnMatrix::Write(AlignedFileWriteStream* fo) const {
std::size_t bytes{0};
bytes += common::WriteVec(fo, index_);
bytes += common::WriteVec(fo, type_);
bytes += common::WriteVec(fo, row_ind_);
bytes += common::WriteVec(fo, feature_offsets_);
bytes += common::WriteVec(fo, missing_.storage);
bytes += fo->Write(bins_type_size_);
bytes += fo->Write(any_missing_);
return bytes;
}
}