Categorical data support in CPU sketching. (#7221)

2021-09-17 04:37:09 +08:00
parent 9f63d6fead
commit 31c1e13f90
7 changed files with 129 additions and 57 deletions
--- a/tests/cpp/common/test_hist_util.cc
+++ b/tests/cpp/common/test_hist_util.cc
@@ -1,3 +1,6 @@
+/*!
+ * Copyright 2019-2021 by XGBoost Contributors
+ */
 #include <gtest/gtest.h>
 #include <vector>
 #include <string>
@@ -388,5 +391,16 @@ TEST(HistUtil, SketchFromWeights) {
  TestSketchFromWeights(true);
  TestSketchFromWeights(false);
 }
+
+TEST(HistUtil, SketchCategoricalFeatures) {
+  TestCategoricalSketch(1000, 256, 32, false,
+                        [](DMatrix *p_fmat, int32_t num_bins) {
+                          return SketchOnDMatrix(p_fmat, num_bins);
+                        });
+  TestCategoricalSketch(1000, 256, 32, true,
+                        [](DMatrix *p_fmat, int32_t num_bins) {
+                          return SketchOnDMatrix(p_fmat, num_bins);
+                        });
+}
 }  // namespace common
 }  // namespace xgboost