Add cudf::explode_position

cpp/include/cudf/reshape.hpp

@@ -139,6 +139,49 @@ std::unique_ptr<table> explode(
   size_type explode_column_idx,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
+/**
+ * @brief Explodes a list column's elements and includes a position column.
+ *
+ * Any list is exploded, which means the elements of the list in each row are expanded into new rows
+ * in the output. The corresponding rows for other columns in the input are duplicated. A position
+ *column is added that has the index inside the original list for each row. Example:
+ * ```
+ * [[5,10,15], 100],
+ * [[20,25],   200],
+ * [[30],      300],
+ * returns
+ * [5,   0,    100],
+ * [10,  1,    100],
+ * [15,  2,    100],
+ * [20,  0,    200],
+ * [25,  1,    200],
+ * [30,  0,    300],
+ * ```
+ *
+ * Nulls and empty lists propagate in different ways depending on what is null or empty.
+ *```
+ * [[5,null,15], 100],
+ * [null,        200],
+ * [[],          300],
+ * returns
+ * [5,    0,     100],
+ * [null, 1,     100],
+ * [15,   2,     100],
+ * ```
+ * Note that null lists are completely removed from the output
+ * and nulls and empty lists inside lists are pulled out and remain.
+ *
+ * @param input_table Table to explode.
+ * @param explode_column_idx Column index to explode inside the table.
+ * @param mr Device memory resource used to allocate the returned column's device memory.
+ *
+ * @return A new table with explode_col exploded.
+ */
+std::unique_ptr<table> pos_explode(
+  table_view const& input_table,
+  size_type explode_column_idx,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
 /** @} */  // end of group
 
 }  // namespace cudf

cpp/src/reshape/explode.cu

@@ -39,77 +39,114 @@ namespace {
 /**
  * @brief Function object for exploding a column.
  */
+template <bool include_pos>
 struct explode_functor {
+  /**
+   * @brief Function object for exploding a column.
+   */
   template <typename T>
-  std::unique_ptr<table> operator()(table_view const& input_table,
-                                    size_type explode_column_idx,
-                                    rmm::cuda_stream_view stream,
-                                    rmm::mr::device_memory_resource* mr) const
+  std::enable_if_t<!std::is_same<T, cudf::list_view>::value, std::unique_ptr<table>> operator()(
+    table_view const& input_table,
+    int const explode_column_idx,
+    rmm::cuda_stream_view stream,
+    rmm::mr::device_memory_resource* mr) const
   {
     CUDF_FAIL("Unsupported non-list column");
 
     return std::make_unique<table>();
   }
-};
 
-template <>
-std::unique_ptr<table> explode_functor::operator()<list_view>(
-  table_view const& input_table,
-  size_type explode_column_idx,
-  rmm::cuda_stream_view stream,
-  rmm::mr::device_memory_resource* mr) const
-{
-  /* we explode by building a gather map that includes the number of entries in each list inside
-   the column for each index. Interestingly, this can be done with lower_bound across the offsets
-   as values between the offsets will all map down to the index below. We have some off-by-one
-   manipulations we need to do with the output, but it's almost our gather map by itself. Once we
-   build the gather map we need to remove the explode column from the table and run gather on it.
-   Next we build the explode column, which turns out is simply lifting the child column out of the
-   explode column. This unrolls the top level of lists. Then we need to insert the explode column
-   back into the table and return it. */
-  lists_column_view lc{input_table.column(explode_column_idx)};
-  auto sliced_child = lc.get_sliced_child(stream);
-  rmm::device_uvector<size_type> gather_map_indices(sliced_child.size(), stream, mr);
-
-  // sliced columns can make this a little tricky. We have to start iterating at the start of the
-  // offsets for this column, which could be > 0. Then we also have to handle rebasing the offsets
-  // as we go.
-  auto offsets           = lc.offsets().begin<size_type>() + lc.offset();
-  auto offsets_minus_one = thrust::make_transform_iterator(
-    offsets, [offsets] __device__(auto i) { return (i - offsets[0]) - 1; });
-  auto counting_iter = thrust::make_counting_iterator(0);
-
-  // This looks like an off-by-one bug, but what is going on here is that we need to reduce each
-  // result from `lower_bound` by 1 to build the correct gather map. It was pointed out that
-  // this can be accomplished by simply skipping the first entry and using the result of
-  // `lower_bound` directly.
-  thrust::lower_bound(rmm::exec_policy(stream),
-                      offsets_minus_one + 1,
-                      offsets_minus_one + lc.size() + 1,
-                      counting_iter,
-                      counting_iter + gather_map_indices.size(),
-                      gather_map_indices.begin());
-
-  auto select_iter = thrust::make_transform_iterator(
-    thrust::make_counting_iterator(0),
-    [explode_column_idx](size_type i) { return i >= explode_column_idx ? i + 1 : i; });
-  std::vector<size_type> selected_columns(select_iter, select_iter + input_table.num_columns() - 1);
-
-  auto gathered_table = cudf::detail::gather(
-    input_table.select(selected_columns),
-    column_view(data_type(type_to_id<size_type>()), sliced_child.size(), gather_map_indices.data()),
-    cudf::out_of_bounds_policy::DONT_CHECK,
-    cudf::detail::negative_index_policy::ALLOWED,
-    stream,
-    mr);
-
-  std::vector<std::unique_ptr<column>> columns = gathered_table.release()->release();
-
-  columns.insert(columns.begin() + explode_column_idx,
-                 std::make_unique<column>(column(sliced_child, stream, mr)));
-
-  return std::make_unique<table>(std::move(columns));
-}
+  /**
+   * @brief Function object for exploding a column.
+   */
+  template <typename T>
+  std::enable_if_t<std::is_same<T, cudf::list_view>::value, std::unique_ptr<table>> operator()(
+    table_view const& input_table,
+    int const explode_column_idx,
+    rmm::cuda_stream_view stream,
+    rmm::mr::device_memory_resource* mr) const
+  {
+    /* we explode by building a gather map that includes the number of entries in each list inside
+     the column for each index. Interestingly, this can be done with lower_bound across the offsets
+     as values between the offsets will all map down to the index below. We have some off-by-one
+     manipulations we need to do with the output, but it's almost our gather map by itself. Once we
+     build the gather map we need to remove the explode column from the table and run gather on it.
+     Next we build the explode column, which turns out is simply lifting the child column out of the
+     explode column. This unrolls the top level of lists. Then we need to insert the explode column
+     back into the table and return it. */
+    lists_column_view lc{input_table.column(explode_column_idx)};
+    auto sliced_child = lc.get_sliced_child(stream);
+    rmm::device_uvector<size_type> gather_map_indices(sliced_child.size(), stream, mr);
+
+    // sliced columns can make this a little tricky. We have to start iterating at the start of the
+    // offsets for this column, which could be > 0. Then we also have to handle rebasing the offsets
+    // as we go.
+    auto offsets           = lc.offsets().begin<size_type>() + lc.offset();
+    auto offsets_minus_one = thrust::make_transform_iterator(
+      offsets, [offsets] __device__(auto i) { return (i - offsets[0]) - 1; });
+    auto counting_iter = thrust::make_counting_iterator(0);
+
+    rmm::device_uvector<size_type> pos(include_pos ? sliced_child.size() : 0, stream, mr);
+
+    // This looks like an off-by-one bug, but what is going on here is that we need to reduce each
+    // result from `lower_bound` by 1 to build the correct gather map. It was pointed out that
+    // this can be accomplished by simply skipping the first entry and using the result of
+    // `lower_bound` directly.
+    if (include_pos) {
+      thrust::transform(rmm::exec_policy(stream),
+                        counting_iter,
+                        counting_iter + gather_map_indices.size(),
+                        gather_map_indices.begin(),
+                        [position_array = pos.data(),
+                         offsets_minus_one,
+                         offsets,
+                         offset_size = lc.size()] __device__(auto idx) -> size_type {
+                          auto lb_idx = thrust::lower_bound(thrust::seq,
+                                                            offsets_minus_one + 1,
+                                                            offsets_minus_one + offset_size + 1,
+                                                            idx) -
+                                        (offsets_minus_one + 1);
+                          position_array[idx] = idx - (offsets[lb_idx] - offsets[0]);
+                          return lb_idx;
+                        });
+    } else {
+      thrust::lower_bound(rmm::exec_policy(stream),
+                          offsets_minus_one + 1,
+                          offsets_minus_one + lc.size() + 1,
+                          counting_iter,
+                          counting_iter + gather_map_indices.size(),
+                          gather_map_indices.begin());
+    }
+
+    auto select_iter = thrust::make_transform_iterator(
+      thrust::make_counting_iterator(0),
+      [explode_column_idx](size_type i) { return i >= explode_column_idx ? i + 1 : i; });
+    std::vector<size_type> selected_columns(select_iter,
+                                            select_iter + input_table.num_columns() - 1);
+
+    auto gathered_table = cudf::detail::gather(
+      input_table.select(selected_columns),
+      column_view(
+        data_type(type_to_id<size_type>()), sliced_child.size(), gather_map_indices.data()),
+      cudf::out_of_bounds_policy::DONT_CHECK,
+      cudf::detail::negative_index_policy::ALLOWED,
+      stream,
+      mr);
+
+    std::vector<std::unique_ptr<column>> columns = gathered_table.release()->release();
+
+    columns.insert(columns.begin() + explode_column_idx,
+                   std::make_unique<column>(column(sliced_child, stream, mr)));
+
+    if (include_pos) {
+      columns.insert(columns.begin() + explode_column_idx,
+                     std::make_unique<column>(column(
+                       data_type(type_to_id<size_type>()), sliced_child.size(), pos.release())));
+    }
+
+    return std::make_unique<table>(std::move(columns));
+  }
+};
 }  // namespace
 
 /**
@@ -118,13 +155,14 @@ std::unique_ptr<table> explode_functor::operator()<list_view>(
  *
  * @param stream CUDA stream used for device memory operations and kernel launches.
  */
+template <bool include_pos>
 std::unique_ptr<table> explode(table_view const& input_table,
                                size_type explode_column_idx,
                                rmm::cuda_stream_view stream,
                                rmm::mr::device_memory_resource* mr)
 {
   return type_dispatcher(input_table.column(explode_column_idx).type(),
-                         explode_functor{},
+                         explode_functor<include_pos>{},
                          input_table,
                          explode_column_idx,
                          stream,
@@ -141,7 +179,18 @@ std::unique_ptr<table> explode(table_view const& input_table,
                                rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::explode(input_table, explode_column_idx, rmm::cuda_stream_default, mr);
+  return detail::explode<false>(input_table, explode_column_idx, rmm::cuda_stream_default, mr);
+}
+
+/**
+ * @copydoc cudf::pos_explode(input_table,explode_column_idx,rmm::mr::device_memory_resource)
+ */
+std::unique_ptr<table> pos_explode(table_view const& input_table,
+                                   size_type explode_column_idx,
+                                   rmm::mr::device_memory_resource* mr)
+{
+  CUDF_FUNC_RANGE();
+  return detail::explode<true>(input_table, explode_column_idx, rmm::cuda_stream_default, mr);
 }
 
 }  // namespace cudf