From 2942dc68e4118aac925110bd78b13454ec936df8 Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Thu, 16 Sep 2021 00:10:25 +0800
Subject: [PATCH] Fix mixed types in GPU sketching. (#7228)

---
 src/common/hist_util.cu            | 32 ++++++++++++++++--------------
 tests/cpp/common/test_hist_util.cu | 26 ++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/src/common/hist_util.cu b/src/common/hist_util.cu
index 1c0961a0e875..33a7c268a6f5 100644
--- a/src/common/hist_util.cu
+++ b/src/common/hist_util.cu
@@ -156,22 +156,24 @@ void RemoveDuplicatedCategories(
   auto d_old_column_sizes_scan = dh::ToSpan(column_sizes_scan);
   dh::caching_device_vector<SketchContainer::OffsetT> new_cuts_size(
       info.num_col_ + 1);
-  auto d_new_cuts_size = dh::ToSpan(new_cuts_size);
-  auto d_new_columns_ptr = dh::ToSpan(new_column_scan);
   CHECK_EQ(new_column_scan.size(), new_cuts_size.size());
-  dh::LaunchN(new_column_scan.size(), [=] __device__(size_t idx) {
-    d_old_column_sizes_scan[idx] = d_new_columns_ptr[idx];
-    if (idx == d_new_columns_ptr.size() - 1) {
-      return;
-    }
-    if (IsCat(d_feature_types, idx)) {
-      // Cut size is the same as number of categories in input.
-      d_new_cuts_size[idx] =
-          d_new_columns_ptr[idx + 1] - d_new_columns_ptr[idx];
-    } else {
-      d_new_cuts_size[idx] = d_cuts_ptr[idx] - d_cuts_ptr[idx];
-    }
-  });
+  dh::LaunchN(
+      new_column_scan.size(),
+      [=, d_new_cuts_size = dh::ToSpan(new_cuts_size),
+       d_old_column_sizes_scan = dh::ToSpan(column_sizes_scan),
+       d_new_columns_ptr = dh::ToSpan(new_column_scan)] __device__(size_t idx) {
+        d_old_column_sizes_scan[idx] = d_new_columns_ptr[idx];
+        if (idx == d_new_columns_ptr.size() - 1) {
+          return;
+        }
+        if (IsCat(d_feature_types, idx)) {
+          // Cut size is the same as number of categories in input.
+          d_new_cuts_size[idx] =
+              d_new_columns_ptr[idx + 1] - d_new_columns_ptr[idx];
+        } else {
+          d_new_cuts_size[idx] = d_cuts_ptr[idx + 1] - d_cuts_ptr[idx];
+        }
+      });
   // Turn size into ptr.
   thrust::exclusive_scan(thrust::device, new_cuts_size.cbegin(),
                          new_cuts_size.cend(), d_cuts_ptr.data());
diff --git a/tests/cpp/common/test_hist_util.cu b/tests/cpp/common/test_hist_util.cu
index a4c18299dfe7..a8a14a8dd948 100644
--- a/tests/cpp/common/test_hist_util.cu
+++ b/tests/cpp/common/test_hist_util.cu
@@ -165,6 +165,32 @@ TEST(HistUtil, DeviceSketchCategoricalFeatures) {
   TestCategoricalSketch(1000, 256, 32, true);
 }
 
+void TestMixedSketch() {
+  size_t n_samples = 1000, n_features = 2, n_categories = 3;
+  std::vector<float> data(n_samples * n_features);
+  SimpleLCG gen;
+  SimpleRealUniformDistribution<float> cat_d{0.0f, float(n_categories)};
+  SimpleRealUniformDistribution<float> num_d{0.0f, 3.0f};
+  for (size_t i = 0; i < n_samples * n_features; ++i) {
+    if (i % 2 == 0) {
+      data[i] = std::floor(cat_d(&gen));
+    } else {
+      data[i] = num_d(&gen);
+    }
+  }
+
+  auto m = GetDMatrixFromData(data, n_samples, n_features);
+  m->Info().feature_types.HostVector().push_back(FeatureType::kCategorical);
+  m->Info().feature_types.HostVector().push_back(FeatureType::kNumerical);
+
+  auto cuts = DeviceSketch(0, m.get(), 64);
+  ASSERT_EQ(cuts.Values().size(), 64 + n_categories);
+}
+
+TEST(HistUtil, DeviceSketchMixedFeatures) {
+  TestMixedSketch();
+}
+
 TEST(HistUtil, DeviceSketchMultipleColumns) {
   int bin_sizes[] = {2, 16, 256, 512};
   int sizes[] = {100, 1000, 1500};