#ifndef XGBOOST_DATA_GRADIENT_INDEX_PAGE_SOURCE_H_
#define XGBOOST_DATA_GRADIENT_INDEX_PAGE_SOURCE_H_
#include <cmath>
#include <cstdint>
#include <memory>
#include <utility>
#include <vector>
#include "../common/hist_util.h"
#include "gradient_index.h"
#include "gradient_index_format.h"
#include "sparse_page_source.h"
#include "xgboost/base.h"
#include "xgboost/data.h"
#include "xgboost/span.h"
namespace xgboost::data {
template <typename S>
class GHistIndexFormatPolicy {
protected:
common::HistogramCuts cuts_;
public:
using FormatT = SparsePageFormat<GHistIndexMatrix>;
public:
[[nodiscard]] auto CreatePageFormat(BatchParam const&) const {
std::unique_ptr<FormatT> fmt{new GHistIndexRawFormat{cuts_}};
return fmt;
}
void SetCuts(common::HistogramCuts cuts) { std::swap(cuts_, cuts); }
};
class GradientIndexPageSource
: public PageSourceIncMixIn<
GHistIndexMatrix, DefaultFormatStreamPolicy<GHistIndexMatrix, GHistIndexFormatPolicy>> {
bool is_dense_;
std::int32_t max_bin_per_feat_;
common::Span<FeatureType const> feature_types_;
double sparse_thresh_;
public:
GradientIndexPageSource(float missing, std::int32_t nthreads, bst_feature_t n_features,
bst_idx_t n_batches, std::shared_ptr<Cache> cache, BatchParam param,
common::HistogramCuts cuts, bool is_dense,
common::Span<FeatureType const> feature_types,
std::shared_ptr<SparsePageSource> source)
: PageSourceIncMixIn(missing, nthreads, n_features, n_batches, cache,
std::isnan(param.sparse_thresh)),
is_dense_{is_dense},
max_bin_per_feat_{param.max_bin},
feature_types_{feature_types},
sparse_thresh_{param.sparse_thresh} {
this->source_ = source;
this->SetCuts(std::move(cuts));
if (this->cuts_.HasCategorical()) {
CHECK(!this->feature_types_.empty());
}
this->Fetch();
}
void Fetch() final;
};
class ExtGradientIndexPageSource
: public ExtQantileSourceMixin<
GHistIndexMatrix, DefaultFormatStreamPolicy<GHistIndexMatrix, GHistIndexFormatPolicy>> {
BatchParam p_;
Context const* ctx_;
DMatrixProxy* proxy_;
MetaInfo* info_;
std::vector<bst_idx_t> base_rows_;
public:
ExtGradientIndexPageSource(
Context const* ctx, float missing, MetaInfo* info, std::shared_ptr<Cache> cache,
BatchParam param, common::HistogramCuts cuts,
std::shared_ptr<DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext>> source,
DMatrixProxy* proxy, std::vector<bst_idx_t> base_rows)
: ExtQantileSourceMixin{missing, ctx->Threads(), static_cast<bst_feature_t>(info->num_col_),
source, cache},
p_{std::move(param)},
ctx_{ctx},
proxy_{proxy},
info_{info},
base_rows_{std::move(base_rows)} {
CHECK(!this->cache_info_->written);
this->source_->Reset();
CHECK(this->source_->Next());
this->SetCuts(std::move(cuts));
this->Fetch();
}
void Fetch() final;
};
} #endif