Reduce quantile sketch safety factor (#12167)

RAMitchell · web-flow · commit 45861e715c0e · 2026-04-17T09:57:15.000+02:00
diff --git a/python-package/xgboost/testing/dask.py b/python-package/xgboost/testing/dask.py
@@ -425,9 +425,8 @@ def run(DMatrixT: Type[dxgb.DaskDMatrix]) -> None:
             evals=[(Xy_valid, "Valid")],
             xgb_model=results["booster"],
         )
-        np.testing.assert_allclose(
-            results_1["history"]["Valid"]["rmse"], results_2["history"]["Valid"]["rmse"]
-        )
+        assert np.isfinite(results_1["history"]["Valid"]["rmse"]).all()
+        assert np.isfinite(results_2["history"]["Valid"]["rmse"]).all()
 
         predt_0 = dxgb.inplace_predict(client, results, denc).compute()
         predt_1 = dxgb.inplace_predict(client, results, dreenc).compute()
diff --git a/src/common/quantile.h b/src/common/quantile.h
@@ -560,7 +560,7 @@ class WQuantileSketch {
   // Safety factor used to oversample the internal sketch relative to the target rank
   // resolution. User-facing epsilon remains the target rank guarantee; `kFactor`
   // only affects how much summary storage we reserve to achieve it.
-  static float constexpr kFactor = 8.0;
+  static float constexpr kFactor = 2.0;
 
  public:
   using Summary = WQSummary<>;
diff --git a/tests/cpp/data/test_gradient_index.cc b/tests/cpp/data/test_gradient_index.cc
@@ -2,25 +2,25 @@
  * Copyright 2021-2024, XGBoost contributors
  */
 #include <gtest/gtest.h>
-#include <xgboost/data.h>                       // for BatchIterator, BatchSet, DMatrix, BatchParam
-
-#include <algorithm>                            // for sort, unique
-#include <cmath>                                // for isnan
-#include <cstddef>                              // for size_t
-#include <limits>                               // for numeric_limits
-#include <memory>                               // for shared_ptr, __shared_ptr_access, unique_ptr
-#include <string>                               // for string
-#include <tuple>                                // for make_tuple, tie, tuple
-#include <utility>                              // for move
-#include <vector>                               // for vector
+#include <xgboost/data.h>  // for BatchIterator, BatchSet, DMatrix, BatchParam
+
+#include <algorithm>  // for sort, unique
+#include <cmath>      // for isnan
+#include <cstddef>    // for size_t
+#include <limits>     // for numeric_limits
+#include <memory>     // for shared_ptr, __shared_ptr_access, unique_ptr
+#include <string>     // for string
+#include <tuple>      // for make_tuple, tie, tuple
+#include <utility>    // for move
+#include <vector>     // for vector
 
 #include "../../../src/common/categorical.h"    // for AsCat
 #include "../../../src/common/column_matrix.h"  // for ColumnMatrix
 #include "../../../src/common/hist_util.h"      // for Index, HistogramCuts, SketchOnDMatrix
-#include "../../../src/common/io.h"             // for MemoryBufferStream
 #include "../../../src/data/adapter.h"          // for SparsePageAdapterBatch
 #include "../../../src/data/gradient_index.h"   // for GHistIndexMatrix
 #include "../../../src/tree/param.h"            // for TrainParam
+#include "../common/test_hist_util.h"           // for ValidateCuts
 #include "../helpers.h"                         // for GenerateRandomCategoricalSingleColumn...
 #include "xgboost/base.h"                       // for bst_bin_t
 #include "xgboost/context.h"                    // for Context
@@ -184,12 +184,8 @@ class GHistIndexMatrixTest : public testing::TestWithParam<std::tuple<float, flo
       ASSERT_EQ(from_sparse_page.Size(), from_ellpack->Size());
       ASSERT_EQ(from_sparse_page.index.Size(), from_ellpack->index.Size());
 
-      auto const &gidx_from_sparse = from_sparse_page.index;
-      auto const &gidx_from_ellpack = from_ellpack->index;
-
-      for (size_t i = 0; i < gidx_from_sparse.Size(); ++i) {
-        ASSERT_EQ(gidx_from_sparse[i], gidx_from_ellpack[i]);
-      }
+      common::ValidateCuts(from_sparse_page.Cuts(), Xy.get(), kBins);
+      common::ValidateCuts(from_ellpack->Cuts(), Xy.get(), kBins);
 
       auto const &columns_from_sparse = from_sparse_page.Transpose();
       auto const &columns_from_ellpack = from_ellpack->Transpose();
@@ -199,20 +195,6 @@ class GHistIndexMatrixTest : public testing::TestWithParam<std::tuple<float, flo
       for (size_t i = 0; i < n_features; ++i) {
         ASSERT_EQ(columns_from_sparse.GetColumnType(i), columns_from_ellpack.GetColumnType(i));
       }
-
-      std::string from_sparse_buf;
-      {
-        common::AlignedMemWriteStream fo{&from_sparse_buf};
-        auto n_bytes = columns_from_sparse.Write(&fo);
-        ASSERT_EQ(fo.Tell(), n_bytes);
-      }
-      std::string from_ellpack_buf;
-      {
-        common::AlignedMemWriteStream fo{&from_ellpack_buf};
-        auto n_bytes = columns_from_sparse.Write(&fo);
-        ASSERT_EQ(fo.Tell(), n_bytes);
-      }
-      ASSERT_EQ(from_sparse_buf, from_ellpack_buf);
     }
   }
 };
diff --git a/tests/cpp/data/test_iterative_dmatrix.cu b/tests/cpp/data/test_iterative_dmatrix.cu
@@ -10,8 +10,9 @@
 #include "../../../src/data/ellpack_page.cuh"
 #include "../../../src/data/ellpack_page.h"
 #include "../../../src/data/iterative_dmatrix.h"
-#include "../../../src/tree/param.h"  // TrainParam
-#include "../filesystem.h"            // for TemporaryDirectory
+#include "../../../src/tree/param.h"   // TrainParam
+#include "../common/test_hist_util.h"  // for ValidateCuts
+#include "../filesystem.h"             // for TemporaryDirectory
 #include "../helpers.h"
 #include "test_iterative_dmatrix.h"
 
@@ -47,35 +48,28 @@ void TestEquivalent(float sparsity) {
 
     std::visit(
         [](auto&& from_iter, auto&& from_data) {
-          ASSERT_EQ(from_iter.gidx_fvalue_map.size(), from_data.gidx_fvalue_map.size());
-          for (size_t i = 0; i < from_iter.gidx_fvalue_map.size(); ++i) {
-            EXPECT_NEAR(from_iter.gidx_fvalue_map[i], from_data.gidx_fvalue_map[i], kRtEps);
-          }
           ASSERT_EQ(from_iter.NumFeatures(), from_data.NumFeatures());
-          for (size_t i = 0; i < from_iter.NumFeatures() + 1; ++i) {
-            ASSERT_EQ(from_iter.feature_segments[i], from_data.feature_segments[i]);
-          }
         },
         from_iter, from_data);
 
+    common::ValidateCuts(page_concatenated->Cuts(), dm.get(), 256);
+    common::ValidateCuts(ellpack.Impl()->Cuts(), dm.get(), 256);
+
     std::vector<common::CompressedByteT> buffer_from_iter, buffer_from_data;
     auto data_iter = page_concatenated->GetHostEllpack(&ctx, &buffer_from_iter);
     auto data_buf = ellpack.Impl()->GetHostEllpack(&ctx, &buffer_from_data);
     ASSERT_NE(buffer_from_data.size(), 0);
     ASSERT_NE(buffer_from_iter.size(), 0);
     CHECK_EQ(ellpack.Impl()->NumSymbols(), page_concatenated->NumSymbols());
 
-    std::visit(
-        [](auto&& from_iter, auto&& from_data) {
-          CHECK_EQ(from_data.n_rows * from_data.row_stride,
-                   from_data.n_rows * from_iter.row_stride);
-        },
-        from_iter, from_data);
+    std::visit([](auto&& from_iter,
+                  auto&& from_data) { CHECK_EQ(from_data.row_stride, from_iter.row_stride); },
+               from_iter, from_data);
     std::visit(
         [](auto&& from_data, auto&& data_buf, auto&& data_iter) {
-          for (size_t i = 0; i < from_data.n_rows * from_data.row_stride; ++i) {
-            CHECK_EQ(data_buf.gidx_iter[i], data_iter.gidx_iter[i]);
-          }
+          ASSERT_EQ(data_buf.row_stride, data_iter.row_stride);
+          ASSERT_EQ(data_buf.NullValue(), data_iter.NullValue());
+          ASSERT_EQ(from_data.n_rows * from_data.row_stride, data_buf.n_rows * data_buf.row_stride);
         },
         from_data, data_buf, data_iter);
   }
diff --git a/tests/test_distributed/test_with_spark/test_spark.py b/tests/test_distributed/test_with_spark/test_spark.py
@@ -245,8 +245,6 @@ def reg_data(self, spark: SparkSession) -> RegData:
     def test_regressor(
         self, spark: SparkSession, reg_data: RegData, num_workers: int
     ) -> None:
-        train_rows = np.where(~reg_data.is_val)[0]
-        validation_rows = np.where(reg_data.is_val)[0]
         device = _spark_test_device(spark)
 
         reg_param = {
@@ -258,13 +256,6 @@ def test_regressor(
             "early_stopping_rounds": 1,
             "device": device,
         }
-        reg = XGBRegressor(**reg_param).fit(
-            reg_data.X_train,
-            reg_data.y_train,
-            sample_weight=reg_data.weights[train_rows],
-            eval_set=[(reg_data.X_test, reg_data.y_test)],
-            sample_weight_eval_set=[reg_data.weights[validation_rows]],
-        )
         spark_regressor = SparkXGBRegressor(
             pred_contrib_col="pred_contribs",
             weight_col="weight",
@@ -285,8 +276,7 @@ def test_regressor(
             .toPandas()["pred_contribs"]
             .tolist()
         )
-        rounds = reg.get_booster().num_boosted_rounds()
-        iter_range = (0, max(1, min(5, rounds)))
+        iter_range = (0, 1)
         spark_iter_regressor = SparkXGBRegressor(
             weight_col="weight",
             validation_indicator_col="is_val",
@@ -302,32 +292,22 @@ def test_regressor(
             .to_numpy()
         )
 
-        score_atol = 1e-2
         train_history = spark_regressor.training_summary.train_objective_history["rmse"]
+        valid_history = spark_regressor.training_summary.validation_objective_history[
+            "rmse"
+        ]
         assert len(train_history) > 0
+        assert len(valid_history) > 0
+        assert len(train_history) == len(valid_history)
         assert np.isfinite(train_history).all()
-        assert np.all(np.diff(train_history) <= 0.0)
-        assert np.allclose(
-            reg.best_score,
-            spark_regressor._xgb_sklearn_model.best_score,
-            atol=score_atol,
-        )
-        assert preds.shape == reg.predict(reg_data.X).shape
-        assert (
-            iter_preds.shape
-            == reg.predict(reg_data.X, iteration_range=iter_range).shape
-        )
+        assert np.isfinite(valid_history).all()
+        assert preds.shape == (len(reg_data.y),)
+        assert iter_preds.shape == preds.shape
 
         assert np.allclose(pred_contribs.sum(axis=1), preds, rtol=1e-3)
         assert np.allclose(
-            reg.evals_result()["validation_0"]["rmse"],
-            spark_regressor.training_summary.validation_objective_history["rmse"],
-            atol=score_atol,
-        )
-        assert np.allclose(
-            reg.best_score,
+            min(valid_history),
             spark_regressor._xgb_sklearn_model.best_score,
-            atol=score_atol,
         )
 
     def test_training_continuation(self, reg_data: RegData) -> None:
