Categorical data support in CPU sketching. (#7221)

tests/cpp/common/test_hist_util.h

@@ -1,3 +1,6 @@
+/*!
+ * Copyright 2019-2021 by XGBoost Contributors
+ */
 #pragma once
 #include <gtest/gtest.h>
 #include <dmlc/filesystem.h>
@@ -5,6 +8,8 @@
 #include <vector>
 #include <string>
 #include <fstream>
+
+#include "../helpers.h"
 #include "../../../src/common/hist_util.h"
 #include "../../../src/data/simple_dmatrix.h"
 #include "../../../src/data/adapter.h"
@@ -206,5 +211,45 @@ inline void ValidateCuts(const HistogramCuts& cuts, DMatrix* dmat,
   }
 }
 
+/**
+ * \brief Test for sketching on categorical data.
+ *
+ * \param sketch Sketch function, can be on device or on host.
+ */
+template <typename Fn>
+void TestCategoricalSketch(size_t n, size_t num_categories, int32_t num_bins,
+                           bool weighted, Fn sketch) {
+  auto x = GenerateRandomCategoricalSingleColumn(n, num_categories);
+  auto dmat = GetDMatrixFromData(x, n, 1);
+  dmat->Info().feature_types.HostVector().push_back(FeatureType::kCategorical);
+
+  if (weighted) {
+    std::vector<float> weights(n, 0);
+    SimpleLCG lcg;
+    SimpleRealUniformDistribution<float> dist(0, 1);
+    for (auto& v : weights) {
+      v = dist(&lcg);
+    }
+    dmat->Info().weights_.HostVector() = weights;
+  }
+
+  ASSERT_EQ(dmat->Info().feature_types.Size(), 1);
+  auto cuts = sketch(dmat.get(), num_bins);
+  std::sort(x.begin(), x.end());
+  auto n_uniques = std::unique(x.begin(), x.end()) - x.begin();
+  ASSERT_NE(n_uniques, x.size());
+  ASSERT_EQ(cuts.TotalBins(), n_uniques);
+  ASSERT_EQ(n_uniques, num_categories);
+
+  auto& values = cuts.cut_values_.HostVector();
+  ASSERT_TRUE(std::is_sorted(values.cbegin(), values.cend()));
+  auto is_unique = (std::unique(values.begin(), values.end()) - values.begin()) == n_uniques;
+  ASSERT_TRUE(is_unique);
+
+  x.resize(n_uniques);
+  for (size_t i = 0; i < n_uniques; ++i) {
+    ASSERT_EQ(x[i], values[i]);
+  }
+}
 }  // namespace common
 }  // namespace xgboost