Convert thrust::transform_iterator<> into pointers

cpp/src/copying/copy.cu

@@ -200,10 +200,13 @@ std::unique_ptr<column> scatter_gather_based_if_else(Left const& lhs,
     auto const gather_lhs = make_counting_transform_iterator(
       size_type{0}, lhs_gather_map_functor<Filter>{is_left, null_map_entry});
 
+    rmm::device_uvector<size_type> gather_map(size, stream);
+    thrust::copy(rmm::exec_policy(stream), gather_lhs, gather_lhs + size, gather_map.begin());
+
     auto const lhs_gathered_columns =
       cudf::detail::gather(table_view{std::vector<cudf::column_view>{lhs}},
-                           gather_lhs,
-                           gather_lhs + size,
+                           gather_map.begin(),
+                           gather_map.end(),
                            out_of_bounds_policy::NULLIFY,
                            stream,
                            mr)

cpp/src/copying/gather.cu

@@ -703,14 +703,16 @@ INSTANTIATE(int16_t*)
 INSTANTIATE(int32_t*)
 INSTANTIATE(int32_t const*)
 INSTANTIATE(int64_t*)
+INSTANTIATE(thrust::constant_iterator<int32_t>)
+INSTANTIATE(thrust::counting_iterator<int32_t>)
 
 INSTANTIATE(uint8_t*)
 INSTANTIATE(uint16_t*)
 INSTANTIATE(uint32_t*)
 INSTANTIATE(uint64_t*)
 
 INSTANTIATE(uint32_t const*)
-
+// INSTANTIATE(thrust::constant_iterator<uint32_t>)
 }  // namespace detail
 }  // namespace cudf
    //

cpp/src/copying/sample.cu

@@ -57,8 +57,11 @@ std::unique_ptr<table> sample(table_view const& input,
     };
 
     auto begin = cudf::detail::make_counting_transform_iterator(0, RandomGen);
+    rmm::device_uvector<size_type> gather_map(n, stream);
+    thrust::copy(rmm::exec_policy(stream), begin, begin + n, gather_map.begin());
 
-    return detail::gather(input, begin, begin + n, out_of_bounds_policy::DONT_CHECK, stream, mr);
+    return detail::gather(
+      input, gather_map.begin(), gather_map.end(), out_of_bounds_policy::DONT_CHECK, stream, mr);
   } else {
     auto gather_map =
       make_numeric_column(data_type{type_id::INT32}, num_rows, mask_state::UNALLOCATED, stream);

cpp/src/filling/repeat.cu

@@ -151,9 +151,16 @@ std::unique_ptr<table> repeat(table_view const& input_table,
   auto output_size = input_table.num_rows() * count;
   auto map_begin   = cudf::detail::make_counting_transform_iterator(
     0, [count] __device__(auto i) { return i / count; });
-  auto map_end = map_begin + output_size;
 
-  return gather(input_table, map_begin, map_end, out_of_bounds_policy::DONT_CHECK, stream, mr);
+  rmm::device_uvector<size_type> gather_map(output_size, stream);
+  thrust::copy(rmm::exec_policy(stream), map_begin, map_begin + output_size, gather_map.begin());
+
+  return gather(input_table,
+                gather_map.begin(),
+                gather_map.end(),
+                out_of_bounds_policy::DONT_CHECK,
+                stream,
+                mr);
 }
 
 }  // namespace detail

cpp/src/groupby/sort/sort_helper.cu

@@ -309,12 +309,12 @@ std::unique_ptr<table> sort_groupby_helper::unique_keys(rmm::cuda_stream_view st
   auto gather_map_it = thrust::make_transform_iterator(
     group_offsets(stream).begin(), [idx_data] __device__(size_type i) { return idx_data[i]; });
 
-  return cudf::detail::gather(_keys,
-                              gather_map_it,
-                              gather_map_it + num_groups(stream),
-                              out_of_bounds_policy::DONT_CHECK,
-                              stream,
-                              mr);
+  auto const N = num_groups(stream);
+  rmm::device_uvector<size_type> gather_map(N, stream);
+  thrust::copy(rmm::exec_policy(stream), gather_map_it, gather_map_it + N, gather_map.begin());
+
+  return cudf::detail::gather(
+    _keys, gather_map.begin(), gather_map.end(), out_of_bounds_policy::DONT_CHECK, stream, mr);
 }
 
 std::unique_ptr<table> sort_groupby_helper::sorted_keys(rmm::cuda_stream_view stream,

cpp/src/lists/copying/gather.cu

@@ -91,9 +91,14 @@ std::unique_ptr<column> gather_list_leaf(column_view const& column,
                                          rmm::mr::device_memory_resource* mr)
 {
   // gather map iterator for this level (N)
-  auto gather_map_begin = thrust::make_transform_iterator(
+  auto const gather_map_iter = thrust::make_transform_iterator(
     thrust::make_counting_iterator<size_type>(0), list_gatherer{gd});
-  size_type gather_map_size = gd.gather_map_size;
+  size_type const gather_map_size = gd.gather_map_size;
+  rmm::device_uvector<size_type> gather_map(gather_map_size, stream);
+  thrust::copy(rmm::exec_policy(stream),
+               gather_map_iter,
+               gather_map_iter + gather_map_size,
+               gather_map.begin());
 
   // call the normal gather
   //  auto leaf_column = cudf::type_dispatcher<dispatch_storage_type>(
@@ -110,8 +115,8 @@ std::unique_ptr<column> gather_list_leaf(column_view const& column,
 
   auto leaf_column =
     std::move(gather(table_view{{column}},
-                     gather_map_begin,
-                     gather_map_begin + gather_map_size,
+                     gather_map.begin(),
+                     gather_map.end(),
                      // note : we don't need to bother checking for out-of-bounds here since
                      // our inputs at this stage aren't coming from the user.
                      out_of_bounds_policy::DONT_CHECK,
@@ -130,8 +135,8 @@ std::unique_ptr<column> gather_list_leaf(column_view const& column,
   if (null_count > 0) {
     auto list_cdv = column_device_view::create(column);
     auto validity = cudf::detail::valid_if(
-      gather_map_begin,
-      gather_map_begin + gd.gather_map_size,
+      gather_map.begin(),
+      gather_map.end(),
       [cdv = *list_cdv] __device__(int index) { return cdv.is_valid(index) ? true : false; },
       stream,
       mr);

cpp/src/lists/copying/segmented_gather.cu

@@ -61,11 +61,16 @@ std::unique_ptr<column> segmented_gather(lists_column_view const& value_column,
     return value_offsets[offset_idx] + wrapped_sub_index - value_offsets[0];
   };
   auto child_gather_index_begin = cudf::detail::make_counting_transform_iterator(0, transformer);
+  rmm::device_uvector<size_type> child_gather_map(gather_map_size, stream);
+  thrust::copy(rmm::exec_policy(stream),
+               child_gather_index_begin,
+               child_gather_index_begin + gather_map_size,
+               child_gather_map.begin());
 
   // Call gather on child of value_column
   auto child_table = cudf::detail::gather(table_view({value_column.get_sliced_child(stream)}),
-                                          child_gather_index_begin,
-                                          child_gather_index_begin + gather_map_size,
+                                          child_gather_map.begin(),
+                                          child_gather_map.end(),
                                           out_of_bounds_policy::DONT_CHECK,
                                           stream,
                                           mr);

cpp/src/partitioning/round_robin.cu

@@ -94,9 +94,12 @@ std::pair<std::unique_ptr<cudf::table>, std::vector<cudf::size_type>> degenerate
     rmm::device_uvector<cudf::size_type> partition_offsets(num_partitions, stream);
     thrust::sequence(rmm::exec_policy(stream), partition_offsets.begin(), partition_offsets.end());
 
+    rmm::device_uvector<cudf::size_type> gather_map(nrows, stream);
+    thrust::copy(
+      rmm::exec_policy(stream), rotated_iter_begin, rotated_iter_begin + nrows, gather_map.begin());
     auto uniq_tbl = cudf::detail::gather(input,
-                                         rotated_iter_begin,
-                                         rotated_iter_begin + nrows,  // map
+                                         gather_map.begin(),
+                                         gather_map.end(),
                                          cudf::out_of_bounds_policy::DONT_CHECK,
                                          stream,
                                          mr);
@@ -225,8 +228,14 @@ std::pair<std::unique_ptr<table>, std::vector<cudf::size_type>> round_robin_part
       return num_partitions * index_within_partition + partition_index;
     });
 
-  auto uniq_tbl = cudf::detail::gather(
-    input, iter_begin, iter_begin + nrows, cudf::out_of_bounds_policy::DONT_CHECK, stream, mr);
+  rmm::device_uvector<size_type> gather_map(nrows, stream);
+  thrust::copy(rmm::exec_policy(stream), iter_begin, iter_begin + nrows, gather_map.begin());
+  auto uniq_tbl = cudf::detail::gather(input,
+                                       gather_map.begin(),
+                                       gather_map.end(),
+                                       cudf::out_of_bounds_policy::DONT_CHECK,
+                                       stream,
+                                       mr);
   auto ret_pair = std::make_pair(std::move(uniq_tbl), std::vector<cudf::size_type>(num_partitions));
 
   // this has the effect of rotating the set of partition sizes

cpp/src/quantiles/quantiles.cu

@@ -48,13 +48,12 @@ std::unique_ptr<table> quantiles(table_view const& input,
   auto const q_device = cudf::detail::make_device_uvector_async(q, stream);
 
   auto quantile_idx_iter = thrust::make_transform_iterator(q_device.begin(), quantile_idx_lookup);
+  rmm::device_uvector<size_type> gather_map(q.size(), stream);
+  thrust::copy(
+    rmm::exec_policy(stream), quantile_idx_iter, quantile_idx_iter + q.size(), gather_map.begin());
 
-  return detail::gather(input,
-                        quantile_idx_iter,
-                        quantile_idx_iter + q.size(),
-                        out_of_bounds_policy::DONT_CHECK,
-                        stream,
-                        mr);
+  return detail::gather(
+    input, gather_map.begin(), gather_map.end(), out_of_bounds_policy::DONT_CHECK, stream, mr);
 }
 
 }  // namespace detail

cpp/src/reshape/tile.cu

@@ -53,9 +53,11 @@ std::unique_ptr<table> tile(const table_view &in,
 
   auto out_num_rows = in_num_rows * count;
   auto tiled_it     = cudf::detail::make_counting_transform_iterator(0, tile_functor{in_num_rows});
+  rmm::device_uvector<size_type> gather_map(out_num_rows, stream);
+  thrust::copy(rmm::exec_policy(stream), tiled_it, tiled_it + out_num_rows, gather_map.begin());
 
   return detail::gather(
-    in, tiled_it, tiled_it + out_num_rows, out_of_bounds_policy::DONT_CHECK, stream, mr);
+    in, gather_map.begin(), gather_map.end(), out_of_bounds_policy::DONT_CHECK, stream, mr);
 }
 }  // namespace detail