Implement scatter for struct columns

cpp/include/cudf/detail/gather.cuh

@@ -466,15 +466,20 @@ struct column_gatherer_impl<struct_view> {
                                                                          mr);
                    });
 
-    gather_bitmask(
-      // Table view of struct column.
-      cudf::table_view{
-        std::vector<cudf::column_view>{structs_column.child_begin(), structs_column.child_end()}},
-      gather_map_begin,
-      output_struct_members,
-      nullify_out_of_bounds ? gather_bitmask_op::NULLIFY : gather_bitmask_op::DONT_CHECK,
-      stream,
-      mr);
+    auto const nullable = std::any_of(structs_column.child_begin(),
+                                      structs_column.child_end(),
+                                      [](auto const& col) { return col.nullable(); });
+    if (nullable) {
+      gather_bitmask(
+        // Table view of struct column.
+        cudf::table_view{
+          std::vector<cudf::column_view>{structs_column.child_begin(), structs_column.child_end()}},
+        gather_map_begin,
+        output_struct_members,
+        nullify_out_of_bounds ? gather_bitmask_op::NULLIFY : gather_bitmask_op::DONT_CHECK,
+        stream,
+        mr);
+    }
 
     return cudf::make_structs_column(
       gather_map_size,
@@ -652,11 +657,15 @@ std::unique_ptr<table> gather(
                                                    mr));
   }
 
-  gather_bitmask_op const op = bounds_policy == out_of_bounds_policy::NULLIFY
-                                 ? gather_bitmask_op::NULLIFY
-                                 : gather_bitmask_op::DONT_CHECK;
-
-  gather_bitmask(source_table, gather_map_begin, destination_columns, op, stream, mr);
+  auto const nullable = bounds_policy == out_of_bounds_policy::NULLIFY ||
+                        std::any_of(source_table.begin(), source_table.end(), [](auto const& col) {
+                          return col.nullable();
+                        });
+  if (nullable) {
+    auto const op = bounds_policy == out_of_bounds_policy::NULLIFY ? gather_bitmask_op::NULLIFY
+                                                                   : gather_bitmask_op::DONT_CHECK;
+    gather_bitmask(source_table, gather_map_begin, destination_columns, op, stream, mr);
+  }
 
   return std::make_unique<table>(std::move(destination_columns));
 }

cpp/include/cudf/detail/scatter.cuh

@@ -25,13 +25,16 @@
 #include <cudf/dictionary/dictionary_column_view.hpp>
 #include <cudf/dictionary/dictionary_factories.hpp>
 #include <cudf/lists/detail/scatter.cuh>
+#include <cudf/null_mask.hpp>
 #include <cudf/strings/detail/scatter.cuh>
 #include <cudf/strings/string_view.cuh>
 #include <cudf/utilities/traits.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
 
+#include <thrust/uninitialized_fill.h>
+
 namespace cudf {
 namespace detail {
 
@@ -42,10 +45,9 @@ namespace detail {
  * function using the PASSTHROUGH op since the resulting map may contain index
  * values outside the target's range.
  *
- * First, the gather-map is initialized with invalid entries.
- * The gather_rows is used since it should always be outside the target size.
- *
- * Then, the `output[scatter_map[i]] = i`.
+ * First, the gather-map is initialized with an invalid index.
+ * The value `numeric_limits::lowest()` is used since it should always be outside the target size.
+ * Then, `output[scatter_map[i]] = i` for each `i`.
  *
  * @tparam MapIterator Iterator type of the input scatter map.
  * @param scatter_map_begin Beginning of scatter map.
@@ -62,11 +64,16 @@ auto scatter_to_gather(MapIterator scatter_map_begin,
 {
   using MapValueType = typename thrust::iterator_traits<MapIterator>::value_type;
 
-  // The gather_map is initialized with gather_rows value to identify pass-through entries
-  // when calling the gather_bitmask() which applies a pass-through whenever it finds a
+  // The gather_map is initialized with `numeric_limits::lowest()` value to identify pass-through
+  // entries when calling the gather_bitmask() which applies a pass-through whenever it finds a
   // value outside the range of the target column.
-  // We'll use the gather_rows value for this since it should always be outside the valid range.
-  auto gather_map = rmm::device_vector<size_type>(gather_rows, gather_rows);
+  // We'll use the `numeric_limits::lowest()` value for this since it should always be outside the
+  // valid range.
+  auto gather_map = rmm::device_uvector<size_type>(gather_rows, stream);
+  thrust::uninitialized_fill(rmm::exec_policy(stream),
+                             gather_map.begin(),
+                             gather_map.end(),
+                             std::numeric_limits<size_type>::lowest());
 
   // Convert scatter map to a gather map
   thrust::scatter(
@@ -79,6 +86,39 @@ auto scatter_to_gather(MapIterator scatter_map_begin,
   return gather_map;
 }
 
+/**
+ * @brief Create a complement map of `scatter_to_gather` map
+ *
+ * The purpose of this map is to create an identity-mapping for the rows that are not
+ * touched by the `scatter_map`.
+ *
+ * The output result of this mapping is firstly initialized as an identity-mapping
+ * (i.e., `output[i] = i`). Then, for each value `idx` from `scatter_map`, the value `output[idx]`
+ * is set to `numeric_limits::lowest()`, which is an invalid, out-of-bound index to identify the
+ * pass-through entries when calling the `gather_bitmask()` function.
+ *
+ */
+template <typename MapIterator>
+auto scatter_to_gather_complement(MapIterator scatter_map_begin,
+                                  MapIterator scatter_map_end,
+                                  size_type gather_rows,
+                                  rmm::cuda_stream_view stream)
+{
+  auto gather_map = rmm::device_uvector<size_type>(gather_rows, stream);
+  thrust::sequence(rmm::exec_policy(stream), gather_map.begin(), gather_map.end(), 0);
+
+  auto const out_of_bounds_begin =
+    thrust::make_constant_iterator(std::numeric_limits<size_type>::lowest());
+  auto const out_of_bounds_end =
+    out_of_bounds_begin + thrust::distance(scatter_map_begin, scatter_map_end);
+  thrust::scatter(rmm::exec_policy(stream),
+                  out_of_bounds_begin,
+                  out_of_bounds_end,
+                  scatter_map_begin,
+                  gather_map.begin());
+  return gather_map;
+}
+
 template <typename Element, typename Enable = void>
 struct column_scatterer_impl {
   std::unique_ptr<column> operator()(...) const { CUDF_FAIL("Unsupported type for scatter."); }
@@ -214,6 +254,89 @@ struct column_scatterer {
   }
 };
 
+template <>
+struct column_scatterer_impl<struct_view> {
+  template <typename MapItRoot>
+  std::unique_ptr<column> operator()(column_view const& source,
+                                     MapItRoot scatter_map_begin,
+                                     MapItRoot scatter_map_end,
+                                     column_view const& target,
+                                     rmm::cuda_stream_view stream,
+                                     rmm::mr::device_memory_resource* mr) const
+  {
+    CUDF_EXPECTS(source.num_children() == target.num_children(),
+                 "Scatter source and target are not of the same type.");
+
+    auto const scatter_map_size = std::distance(scatter_map_begin, scatter_map_end);
+    if (scatter_map_size == 0) { return std::make_unique<column>(target, stream, mr); }
+
+    structs_column_view const structs_src(source);
+    structs_column_view const structs_target(target);
+    std::vector<std::unique_ptr<column>> output_struct_members(structs_src.num_children());
+
+    std::transform(structs_src.child_begin(),
+                   structs_src.child_end(),
+                   structs_target.child_begin(),
+                   output_struct_members.begin(),
+                   [&scatter_map_begin, &scatter_map_end, stream, mr](auto const& source_col,
+                                                                      auto const& target_col) {
+                     return type_dispatcher<dispatch_storage_type>(source_col.type(),
+                                                                   column_scatterer{},
+                                                                   source_col,
+                                                                   scatter_map_begin,
+                                                                   scatter_map_end,
+                                                                   target_col,
+                                                                   stream,
+                                                                   mr);
+                   });
+
+    // We still need to call `gather_bitmask` even when the source's children are not nullable,
+    // as if the target's children have null_masks, those null_masks need to be updated after
+    // being scattered onto
+    auto const child_nullable = std::any_of(structs_src.child_begin(),
+                                            structs_src.child_end(),
+                                            [](auto const& col) { return col.nullable(); }) or
+                                std::any_of(structs_target.child_begin(),
+                                            structs_target.child_end(),
+                                            [](auto const& col) { return col.nullable(); });
+    if (child_nullable) {
+      auto const gather_map =
+        scatter_to_gather(scatter_map_begin, scatter_map_end, source.size(), stream);
+      gather_bitmask(cudf::table_view{std::vector<cudf::column_view>{structs_src.child_begin(),
+                                                                     structs_src.child_end()}},
+                     gather_map.begin(),
+                     output_struct_members,
+                     gather_bitmask_op::PASSTHROUGH,
+                     stream,
+                     mr);
+    }
+
+    // Need to put the result column in a vector to call `gather_bitmask`
+    std::vector<std::unique_ptr<column>> result;
+    result.emplace_back(cudf::make_structs_column(source.size(),
+                                                  std::move(output_struct_members),
+                                                  0,
+                                                  rmm::device_buffer{0, stream, mr},
+                                                  stream,
+                                                  mr));
+
+    // Only gather bitmask from the target column for the rows that have not been scattered onto
+    // The bitmask from the source column will be gathered at the top level `scatter()` call
+    if (target.nullable()) {
+      auto const gather_map =
+        scatter_to_gather_complement(scatter_map_begin, scatter_map_end, target.size(), stream);
+      gather_bitmask(table_view{std::vector<cudf::column_view>{target}},
+                     gather_map.begin(),
+                     result,
+                     gather_bitmask_op::PASSTHROUGH,
+                     stream,
+                     mr);
+    }
+
+    return std::move(result.front());
+  }
+};
+
 /**
  * @brief Scatters the rows of the source table into a copy of the target table
  * according to a scatter map.
@@ -278,10 +401,8 @@ std::unique_ptr<table> scatter(
   // Transform negative indices to index + target size
   auto updated_scatter_map_begin =
     thrust::make_transform_iterator(scatter_map_begin, index_converter<MapType>{target.num_rows()});
-
   auto updated_scatter_map_end =
     thrust::make_transform_iterator(scatter_map_end, index_converter<MapType>{target.num_rows()});
-
   auto result = std::vector<std::unique_ptr<column>>(target.num_columns());
 
   std::transform(source.begin(),
@@ -299,11 +420,16 @@ std::unique_ptr<table> scatter(
                                                                  mr);
                  });
 
-  auto gather_map = scatter_to_gather(
-    updated_scatter_map_begin, updated_scatter_map_end, target.num_rows(), stream);
-
-  gather_bitmask(source, gather_map.begin(), result, gather_bitmask_op::PASSTHROUGH, stream, mr);
-
+  // We still need to call `gather_bitmask` even when the source columns are not nullable,
+  // as if the target has null_mask, that null_mask needs to be updated after scattering
+  auto const nullable =
+    std::any_of(source.begin(), source.end(), [](auto const& col) { return col.nullable(); }) or
+    std::any_of(target.begin(), target.end(), [](auto const& col) { return col.nullable(); });
+  if (nullable) {
+    auto const gather_map = scatter_to_gather(
+      updated_scatter_map_begin, updated_scatter_map_end, target.num_rows(), stream);
+    gather_bitmask(source, gather_map.begin(), result, gather_bitmask_op::PASSTHROUGH, stream, mr);
+  }
   return std::make_unique<table>(std::move(result));
 }
 }  // namespace detail

cpp/tests/CMakeLists.txt

@@ -201,24 +201,25 @@ ConfigureTest(SORT_TEST
 ###################################################################################################
 # - copying tests ---------------------------------------------------------------------------------
 ConfigureTest(COPYING_TEST
-    copying/utility_tests.cpp
+    copying/concatenate_tests.cu
+    copying/copy_range_tests.cpp
+    copying/copy_tests.cu    
+    copying/detail_gather_tests.cu
+    copying/gather_struct_tests.cu        
     copying/gather_tests.cu
     copying/gather_str_tests.cu
     copying/gather_list_tests.cu
-    copying/segmented_gather_list_tests.cpp
-    copying/gather_struct_tests.cu
-    copying/detail_gather_tests.cu
+    copying/get_value_tests.cpp
+    copying/pack_tests.cu
+    copying/sample_tests.cpp
     copying/scatter_tests.cpp
     copying/scatter_list_tests.cu
-    copying/copy_range_tests.cpp
+    copying/scatter_struct_tests.cu
+    copying/segmented_gather_list_tests.cpp
+    copying/shift_tests.cpp
     copying/slice_tests.cpp
     copying/split_tests.cpp
-    copying/copy_tests.cu
-    copying/shift_tests.cpp
-    copying/get_value_tests.cpp
-    copying/sample_tests.cpp
-    copying/concatenate_tests.cu
-    copying/pack_tests.cu)
+    copying/utility_tests.cpp)
 
 ###################################################################################################
 # - utilities tests -------------------------------------------------------------------------------