Update collect_comm to handle value of tuple type (#4410)

Update collect_comm to handle value of tuple type

Authors:
  - Naim (https://github.com/naimnv)

Approvers:
  - Chuck Hastings (https://github.com/ChuckHastings)
  - Seunghwa Kang (https://github.com/seunghwak)
  - Joseph Nke (https://github.com/jnke2016)

URL: #4410

Author: naimnv

cpp/src/utilities/collect_comm.cuh

@@ -100,7 +100,10 @@ collect_values_for_keys(raft::handle_t const& handle,
 
     auto rx_values_for_unique_keys = allocate_dataframe_buffer<value_t>(0, handle.get_stream());
     std::tie(rx_values_for_unique_keys, std::ignore) =
-      shuffle_values(comm, values_for_rx_unique_keys.begin(), rx_value_counts, handle.get_stream());
+      shuffle_values(comm,
+                     get_dataframe_buffer_begin(values_for_rx_unique_keys),
+                     rx_value_counts,
+                     handle.get_stream());
 
     values_for_unique_keys = std::move(rx_values_for_unique_keys);
   }
@@ -136,9 +139,9 @@ collect_values_for_keys(raft::handle_t const& handle,
                                         handle.get_stream());
 
     unique_keys.resize(0, handle.get_stream());
-    values_for_unique_keys.resize(0, handle.get_stream());
+    resize_dataframe_buffer(values_for_unique_keys, 0, handle.get_stream());
     unique_keys.shrink_to_fit(handle.get_stream());
-    values_for_unique_keys.shrink_to_fit(handle.get_stream());
+    shrink_to_fit_dataframe_buffer(values_for_unique_keys, handle.get_stream());
   }
   auto unique_key_value_store_view = unique_key_value_store.view();
 
@@ -248,15 +251,15 @@ collect_values_for_unique_int_vertices(raft::handle_t const& handle,
   thrust::transform(handle.get_thrust_policy(),
                     rx_int_vertices.begin(),
                     rx_int_vertices.end(),
-                    value_buffer.begin(),
+                    get_dataframe_buffer_begin(value_buffer),
                     [local_value_first, local_int_vertex_first] __device__(auto v) {
                       return local_value_first[v - local_int_vertex_first];
                     });
 
   // 3: Shuffle results back to original GPU
 
-  std::tie(value_buffer, std::ignore) =
-    shuffle_values(comm, value_buffer.begin(), rx_int_vertex_counts, handle.get_stream());
+  std::tie(value_buffer, std::ignore) = shuffle_values(
+    comm, get_dataframe_buffer_begin(value_buffer), rx_int_vertex_counts, handle.get_stream());
 
   return std::make_tuple(std::move(collect_unique_int_vertices), std::move(value_buffer));
 }
@@ -305,7 +308,7 @@ collect_values_for_int_vertices(
   thrust::transform(handle.get_thrust_policy(),
                     collect_vertex_first,
                     collect_vertex_last,
-                    value_buffer.begin(),
+                    get_dataframe_buffer_begin(value_buffer),
                     [device_view] __device__(auto v) { return device_view.find(v); });
 
   return value_buffer;