From d3ba8bec07de47dece03f915af7a64c8f446e1b5 Mon Sep 17 00:00:00 2001
From: David Wendt <dwendt@nvidia.com>
Date: Fri, 8 Mar 2024 17:59:35 -0500
Subject: [PATCH] Fix cudf::test::to_host return of host_vector

---
 cpp/include/cudf/detail/gather.cuh         |  2 --
 cpp/include/cudf_test/column_utilities.hpp |  6 ++--
 cpp/tests/utilities/column_utilities.cu    | 32 ++++++++++------------
 3 files changed, 17 insertions(+), 23 deletions(-)

diff --git a/cpp/include/cudf/detail/gather.cuh b/cpp/include/cudf/detail/gather.cuh
index 311a100a21b..6492aa23e80 100644
--- a/cpp/include/cudf/detail/gather.cuh
+++ b/cpp/include/cudf/detail/gather.cuh
@@ -41,9 +41,7 @@
 
 #include <thrust/functional.h>
 #include <thrust/gather.h>
-#include <thrust/host_vector.h>
 #include <thrust/iterator/counting_iterator.h>
-#include <thrust/iterator/transform_iterator.h>
 #include <thrust/logical.h>
 
 #include <algorithm>
diff --git a/cpp/include/cudf_test/column_utilities.hpp b/cpp/include/cudf_test/column_utilities.hpp
index cbfd7a5e45c..a8957473175 100644
--- a/cpp/include/cudf_test/column_utilities.hpp
+++ b/cpp/include/cudf_test/column_utilities.hpp
@@ -174,9 +174,9 @@ bool validate_host_masks(std::vector<bitmask_type> const& expected_mask,
 template <typename T, std::enable_if_t<not cudf::is_fixed_point<T>()>* = nullptr>
 std::pair<thrust::host_vector<T>, std::vector<bitmask_type>> to_host(column_view c)
 {
-  thrust::host_vector<T> host_data(c.size());
-  CUDF_CUDA_TRY(cudaMemcpy(host_data.data(), c.data<T>(), c.size() * sizeof(T), cudaMemcpyDefault));
-  return {host_data, bitmask_to_host(c)};
+  auto col_span  = cudf::device_span<T const>(c.data<T>(), c.size());
+  auto host_data = cudf::detail::make_host_vector_sync(col_span, cudf::get_default_stream());
+  return {std::move(host_data), bitmask_to_host(c)};
 }
 
 // This signature is identical to the above overload apart from SFINAE so
diff --git a/cpp/tests/utilities/column_utilities.cu b/cpp/tests/utilities/column_utilities.cu
index a556a8702bd..2cd7dc1574c 100644
--- a/cpp/tests/utilities/column_utilities.cu
+++ b/cpp/tests/utilities/column_utilities.cu
@@ -906,20 +906,18 @@ void expect_column_empty(cudf::column_view const& col)
 std::vector<bitmask_type> bitmask_to_host(cudf::column_view const& c)
 {
   if (c.nullable()) {
-    auto num_bitmasks = num_bitmask_words(c.size());
-    std::vector<bitmask_type> host_bitmask(num_bitmasks);
-    if (c.offset() == 0) {
-      CUDF_CUDA_TRY(cudaMemcpy(host_bitmask.data(),
-                               c.null_mask(),
-                               num_bitmasks * sizeof(bitmask_type),
-                               cudaMemcpyDefault));
-    } else {
+    auto num_bitmasks      = num_bitmask_words(c.size());
+    auto [bitmask_span, _] = [&] {
+      if (c.offset() == 0) {
+        return std::pair{cudf::device_span<bitmask_type const>(c.null_mask(), num_bitmasks),
+                         rmm::device_buffer{}};
+      }
       auto mask = copy_bitmask(c.null_mask(), c.offset(), c.offset() + c.size());
-      CUDF_CUDA_TRY(cudaMemcpy(
-        host_bitmask.data(), mask.data(), num_bitmasks * sizeof(bitmask_type), cudaMemcpyDefault));
-    }
-
-    return host_bitmask;
+      return std::pair{cudf::device_span<bitmask_type const>(
+                         static_cast<bitmask_type*>(mask.data()), num_bitmasks),
+                       std::move(mask)};
+    }();
+    return cudf::detail::make_std_vector_sync(bitmask_span, cudf::get_default_stream());
   } else {
     return std::vector<bitmask_type>{};
   }
@@ -946,16 +944,14 @@ std::pair<thrust::host_vector<T>, std::vector<bitmask_type>> to_host(column_view
   using namespace numeric;
   using Rep = typename T::rep;
 
-  auto host_rep_types = thrust::host_vector<Rep>(c.size());
-
-  CUDF_CUDA_TRY(
-    cudaMemcpy(host_rep_types.data(), c.begin<Rep>(), c.size() * sizeof(Rep), cudaMemcpyDefault));
+  auto col_span       = cudf::device_span<Rep const>(c.begin<Rep>(), c.size());
+  auto host_rep_types = cudf::detail::make_host_vector_sync(col_span, cudf::get_default_stream());
 
   auto to_fp = [&](Rep val) { return T{scaled_integer<Rep>{val, scale_type{c.type().scale()}}}; };
   auto begin = thrust::make_transform_iterator(std::cbegin(host_rep_types), to_fp);
   auto const host_fixed_points = thrust::host_vector<T>(begin, begin + c.size());
 
-  return {host_fixed_points, bitmask_to_host(c)};
+  return {std::move(host_fixed_points), bitmask_to_host(c)};
 }
 
 template std::pair<thrust::host_vector<numeric::decimal32>, std::vector<bitmask_type>> to_host(