Segmented apply_boolean_mask for LIST columns

cpp/CMakeLists.txt

@@ -359,6 +359,7 @@ add_library(
   src/join/mixed_join_size_kernel_nulls.cu
   src/join/mixed_join_size_kernels_semi.cu
   src/join/semi_join.cu
+  src/lists/apply_boolean_mask.cu
   src/lists/contains.cu
   src/lists/combine/concatenate_list_elements.cu
   src/lists/combine/concatenate_rows.cu

cpp/include/cudf/lists/detail/stream_compaction.hpp

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cudf/column/column.hpp>
+#include <cudf/lists/lists_column_view.hpp>
+
+#include <rmm/mr/device/device_memory_resource.hpp>
+
+namespace cudf::lists::detail {
+
+/**
+ * @copydoc cudf::lists::apply_boolean_mask(lists_column_view const&, lists_column_view const&,
+ * rmm::mr::device_memory_resource*)
+ *
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ */
+std::unique_ptr<column> apply_boolean_mask(
+  lists_column_view const& input,
+  lists_column_view const& boolean_mask,
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+}  // namespace cudf::lists::detail

cpp/include/cudf/lists/stream_compaction.hpp

@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cudf/column/column.hpp>
+#include <cudf/lists/lists_column_view.hpp>
+
+#include <rmm/mr/device/device_memory_resource.hpp>
+
+namespace cudf::lists {
+
+/**
+ * @brief Filters elements in each row of `input` LIST column using `boolean_mask`
+ * LIST of booleans as a mask.
+ *
+ * Given an input `LIST` column and a list-of-bools column, the function produces
+ * a new `LIST` column of the same type as `input`, where each element is copied
+ * from the input row *only* if the corresponding `boolean_mask` is non-null and `true`.
+ *
+ * E.g.
+ * @code{.pseudo}
+ * input        = { {0,1,2}, {3,4}, {5,6,7}, {8,9} };
+ * boolean_mask = { {0,1,1}, {1,0}, {1,1,1}, {0,0} };
+ * results      = { {1,2},   {3},   {5,6,7}, {} };
+ * @endcode
+ *
+ * `input` and `boolean_mask` must have the same number of rows.
+ * The output column has the same number of rows as the input column.
+ * An element is copied to an output row *only* if the corresponding boolean_mask element is `true`.
+ * An output row is invalid only if the input row is invalid.
+ *
+ * @throws cudf::logic_error if `boolean_mask` is not a "lists of bools" column
+ * @throws cudf::logic_error if `input` and `boolean_mask` have different number of rows
+ *
+ * @param input The input list column view to be filtered
+ * @param boolean_mask A nullable list of bools column used to filter `input` elements
+ * @param mr Device memory resource used to allocate the returned table's device memory
+ * @return List column of the same type as `input`, containing filtered list rows
+ */
+std::unique_ptr<column> apply_boolean_mask(
+  lists_column_view const& input,
+  lists_column_view const& boolean_mask,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+}  // namespace cudf::lists

cpp/src/lists/apply_boolean_mask.cu

@@ -0,0 +1,105 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cudf/column/column_factories.hpp>
+#include <cudf/detail/copy.hpp>
+#include <cudf/detail/fill.hpp>
+#include <cudf/detail/iterator.cuh>
+#include <cudf/detail/null_mask.hpp>
+#include <cudf/detail/reduction_functions.hpp>
+#include <cudf/detail/replace.hpp>
+#include <cudf/detail/stream_compaction.hpp>
+#include <cudf/lists/detail/stream_compaction.hpp>
+#include <cudf/lists/stream_compaction.hpp>
+#include <cudf/utilities/bit.hpp>
+
+#include <rmm/exec_policy.hpp>
+
+#include <thrust/reduce.h>
+
+namespace cudf::lists {
+namespace detail {
+
+std::unique_ptr<column> apply_boolean_mask(lists_column_view const& input,
+                                           lists_column_view const& boolean_mask,
+                                           rmm::cuda_stream_view stream,
+                                           rmm::mr::device_memory_resource* mr)
+{
+  CUDF_EXPECTS(boolean_mask.child().type().id() == type_id::BOOL8, "Mask must be of type BOOL8.");
+  CUDF_EXPECTS(input.size() == boolean_mask.size(),
+               "Boolean masks column must have same number of rows as input.");
+  auto const num_rows = input.size();
+
+  if (num_rows == 0) { return cudf::empty_like(input.parent()); }
+
+  auto constexpr offset_data_type = data_type{type_id::INT32};
+
+  auto const boolean_mask_sliced_child = boolean_mask.get_sliced_child(stream);
+
+  auto const make_filtered_child = [&] {
+    auto filtered =
+      cudf::detail::apply_boolean_mask(
+        cudf::table_view{{input.get_sliced_child(stream)}}, boolean_mask_sliced_child, stream, mr)
+        ->release();
+    return std::move(filtered.front());
+  };
+
+  auto const make_output_offsets = [&] {
+    auto boolean_mask_sliced_offsets =
+      cudf::detail::slice(
+        boolean_mask.offsets(), {boolean_mask.offset(), boolean_mask.size() + 1}, stream)
+        .front();
+    auto const sizes       = cudf::reduction::segmented_sum(boolean_mask_sliced_child,
+                                                      boolean_mask_sliced_offsets,
+                                                      offset_data_type,
+                                                      null_policy::EXCLUDE,
+                                                      stream);
+    auto const d_sizes     = column_device_view::create(*sizes, stream);
+    auto const sizes_begin = cudf::detail::make_null_replacement_iterator(*d_sizes, offset_type{0});
+    auto const sizes_end   = sizes_begin + sizes->size();
+    auto output_offsets    = cudf::make_numeric_column(
+      offset_data_type, num_rows + 1, mask_state::UNALLOCATED, stream, mr);
+    auto output_offsets_view = output_offsets->mutable_view();
+
+    // Could have attempted an exclusive_scan(), but it would not compute the last entry.
+    // Instead, inclusive_scan(), followed by writing `0` to the head of the offsets column.
+    thrust::inclusive_scan(rmm::exec_policy(stream),
+                           sizes_begin,
+                           sizes_end,
+                           output_offsets_view.begin<offset_type>() + 1);
+    CUDF_CUDA_TRY(cudaMemsetAsync(
+      output_offsets_view.begin<offset_type>(), 0, sizeof(offset_type), stream.value()));
+    return output_offsets;
+  };
+
+  return cudf::make_lists_column(input.size(),
+                                 make_output_offsets(),
+                                 make_filtered_child(),
+                                 input.null_count(),
+                                 cudf::detail::copy_bitmask(input.parent(), stream, mr),
+                                 stream,
+                                 mr);
+}
+}  // namespace detail
+
+std::unique_ptr<column> apply_boolean_mask(lists_column_view const& input,
+                                           lists_column_view const& boolean_mask,
+                                           rmm::mr::device_memory_resource* mr)
+{
+  return detail::apply_boolean_mask(input, boolean_mask, rmm::cuda_stream_default, mr);
+}
+
+}  // namespace cudf::lists

cpp/tests/CMakeLists.txt

@@ -469,6 +469,7 @@ ConfigureTest(AST_TEST ast/transform_tests.cpp)
 # * lists tests ----------------------------------------------------------------------------------
 ConfigureTest(
   LISTS_TEST
+  lists/apply_boolean_mask_test.cpp
   lists/combine/concatenate_list_elements_tests.cpp
   lists/combine/concatenate_rows_tests.cpp
   lists/contains_tests.cpp