Clean up host quantile helpers#12074
Conversation
# Conflicts: # src/common/quantile.cc
There was a problem hiding this comment.
Pull request overview
This PR refactors the CPU quantile sketching implementation by consolidating the sorted/unsorted sketch containers into a single HostSketchContainer, and updates unit tests accordingly (including moving some previously-header-only helpers into test code).
Changes:
- Remove
SortedSketchContainer/SketchContainerImplsplit and useHostSketchContainerfor both row-page and sorted column-page sketching. - Move group-weight “unroll” helper implementation out of the public header into the
.cc, and add a test-only equivalent helper. - Update/clean up affected tests (including removing the
SearchGroupIndFromRowtest after the method removal).
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/cpp/common/test_quantile.cc | Updates distributed quantile tests to always use HostSketchContainer and template-dispatch PushRowPage vs PushColPage. |
| tests/cpp/common/test_hist_util.h | Adds a test helper for unrolling group weights and switches validation to use it. |
| tests/cpp/common/test_hist_util.cc | Removes test coverage for SearchGroupIndFromRow after its removal. |
| src/common/quantile.h | Renames/reshapes the sketch container API to HostSketchContainer and removes the sorted-container class. |
| src/common/quantile.cc | Moves/implements sketch container logic under HostSketchContainer, adds local UnrollGroupWeights, and refactors cut/category helpers. |
| src/common/hist_util.cc | Updates sorted sketching path to instantiate HostSketchContainer instead of SortedSketchContainer. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| size_t cur_group = 0; | ||
| for (bst_idx_t i = 0; i < n_samples; ++i) { | ||
| results[i] = group_weights[cur_group]; | ||
| if (i == group_ptr[cur_group + 1]) { | ||
| cur_group++; | ||
| } |
There was a problem hiding this comment.
UnrollGroupWeights assigns group_weights[cur_group] to row i before advancing cur_group when i reaches the next group boundary (group_ptr[cur_group + 1]). Since group_ptr_ is a prefix-sum with half-open intervals [group_ptr[g], group_ptr[g+1]), this is off by one and will apply the previous group's weight to the first row of each subsequent group. Adjust the boundary update logic (and the similar loop in MergeWeights) so the group index is advanced before emitting the weight for rows at i == group_ptr[cur_group + 1].
| out[i] = group_weights[cur_group]; | ||
| if (i == group_ptr[cur_group + 1]) { | ||
| cur_group++; | ||
| } |
There was a problem hiding this comment.
The group-boundary handling in UnrollGroupWeightsForTest is off by one: it writes the current group's weight for row i and only then increments cur_group when i == group_ptr[cur_group + 1]. With the usual [group_ptr[g], group_ptr[g+1]) convention, the row at group_ptr[g+1] belongs to the next group, so this will mis-assign weights. Update the loop to advance the group index before assigning the weight for rows at the boundary.
| out[i] = group_weights[cur_group]; | |
| if (i == group_ptr[cur_group + 1]) { | |
| cur_group++; | |
| } | |
| if (i >= group_ptr[cur_group + 1]) { | |
| ++cur_group; | |
| } | |
| out[i] = group_weights[cur_group]; |
There was a problem hiding this comment.
This bug has existed since about 2022 and meant that in the ranking case with weights the quantiles were probably wrong. I fixed and tested it here.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| ParallelFor(sketches_.size(), n_threads_, Sched::Auto(), [&](auto i) { | ||
| auto n_bins = std::min(static_cast<bst_idx_t>(max_bins_), columns_size_[i]); | ||
| n_bins = std::max(n_bins, static_cast<decltype(n_bins)>(1)); | ||
| auto eps = 1.0 / (static_cast<float>(n_bins) * WQSketch::kFactor); | ||
| if (!IsCat(this->feature_types_, i)) { | ||
| sketches_[i] = WQSketch{columns_size_[i], eps}; | ||
| } |
There was a problem hiding this comment.
HostSketchContainer now initializes WQuantileSketch with eps = 1 / (n_bins * kFactor) for all use cases. Previously the sorted-column path used SortedSketchContainer with a much larger eps (e.g. 2.0 / max_bins), which kept WQuantileSketch::limit_size_ relatively small.
With the smaller eps, limit_size_ can become ~O(max_bins * kFactor * log(maxn)) for large columns, and WQuantileSketch::PushSummary() will Reserve(limit_size_ * 2) even when PushColPage() uses PushSorted() (where the incoming summary is already pruned to ~max_bins). On wide datasets this can significantly increase per-feature memory and CPU for the sorted sketch path.
Consider restoring a separate initialization policy for the sorted-column sketching case (e.g. pass a flag into HostSketchContainer or reintroduce a lightweight sorted-only container) so that the sorted path uses a larger eps / smaller limit_size_ consistent with the previous behavior, while keeping the current initialization for the row-wise path.
There was a problem hiding this comment.
This is intentional. I don't see a reason to have less accurate quantiles for the sorted version. The increase is only a small constant factor.
| info.weights_.HostVector() = {1.0f, 5.0f, 9.0f}; | ||
|
|
||
| std::vector<float> expected{1.0f, 1.0f, 5.0f, 9.0f, 9.0f, 9.0f}; | ||
| ASSERT_EQ(UnrollGroupWeightsForTest(info), expected); |
There was a problem hiding this comment.
Why checking a test utility function instead of the real one?
There was a problem hiding this comment.
Sorry that was an artefact from the refactor
Summary
quantile.hTesting