Preserve exact weighted values in sorted sketch (#12148)

Author: RAMitchell

src/common/quantile.h

@@ -118,17 +118,54 @@ struct WQSummary {
     auto const *col_data = column.data();
     auto const col_size = column.size();
     double sum_total{0.0};
+    std::size_t unique_values{0};
     double rmin{0.0};
     double wmin{0.0};
     bst_float last_fvalue{0.0f};
     double next_goal{-1.0f};
 
     // first pass
     for (size_t i = 0; i < col_size; ++i) {
+      if (i == 0 || col_data[i - 1].fvalue != col_data[i].fvalue) {
+        ++unique_values;
+      }
       auto const &c = col_data[i];
       sum_total += weights[c.index];
     }
 
+    if (unique_values <= max_size) {
+      // When we have enough budget to keep every unique feature value, emit the exact
+      // weighted summary instead of running the weighted goal-selection logic below.
+      for (size_t i = 0; i < col_size; ++i) {
+        auto const &c = col_data[i];
+        if (i == 0) {
+          last_fvalue = c.fvalue;
+          wmin = weights[c.index];
+          continue;
+        }
+        if (last_fvalue == c.fvalue) {
+          wmin += weights[c.index];
+          continue;
+        }
+
+        auto rmax = rmin + wmin;
+        data_[this->Size()] = Entry(static_cast<bst_float>(rmin), static_cast<bst_float>(rmax),
+                                    static_cast<bst_float>(wmin), last_fvalue);
+        this->SetSize(this->Size() + 1);
+        rmin = rmax;
+        last_fvalue = c.fvalue;
+        wmin = weights[c.index];
+      }
+
+      if (col_size != 0) {
+        auto rmax = rmin + wmin;
+        data_[this->Size()] = Entry(static_cast<bst_float>(rmin), static_cast<bst_float>(rmax),
+                                    static_cast<bst_float>(wmin), last_fvalue);
+        this->SetSize(this->Size() + 1);
+      }
+      return;
+    }
+
     // second pass
     for (size_t i = 0; i < col_size; ++i) {
       auto const &c = col_data[i];

tests/cpp/common/test_hist_util.cc

@@ -179,6 +179,27 @@ TEST(HistUtil, DenseCutsAccuracyTestWeights) {
   }
 }
 
+TEST(HistUtil, SortedWeightedExactCuts) {
+  Context ctx;
+  std::vector<float> x{0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
+  std::vector<float> weights{3.0f, 0.25f, 7.0f, 2.0f, 0.5f, 4.0f};
+  auto dmat = GetDMatrixFromData(x, x.size(), 1);
+  dmat->Info().weights_.HostVector() = weights;
+
+  auto sorted_cuts = SketchOnDMatrix(&ctx, dmat.get(), x.size(), true);
+  auto row_cuts = SketchOnDMatrix(&ctx, dmat.get(), x.size(), false);
+
+  ASSERT_EQ(sorted_cuts.Ptrs(), row_cuts.Ptrs());
+  ASSERT_EQ(sorted_cuts.Values().size(), row_cuts.Values().size());
+  ASSERT_EQ(sorted_cuts.Values().size(), x.size());
+  for (std::size_t i = 1; i < x.size(); ++i) {
+    EXPECT_FLOAT_EQ(sorted_cuts.Values()[i - 1], x[i]);
+    EXPECT_FLOAT_EQ(sorted_cuts.Values()[i - 1], row_cuts.Values()[i - 1]);
+  }
+  EXPECT_GT(sorted_cuts.Values().back(), x.back());
+  EXPECT_FLOAT_EQ(sorted_cuts.Values().back(), row_cuts.Values().back());
+}
+
 void TestQuantileWithHessian(bool use_sorted) {
   int bin_sizes[] = {2, 16, 256, 512};
   int sizes[] = {1000, 1500};